from numpy import *
import matplotlib.pyplot as plt

# Define any paths here to make things easier.
# A file path simply becomes HOME + 'file', for example.
# note: trailing backslash.

HOME = 'C:\Users\this_user\Desktop\\'

# RIAA as an array of poles and zeros, Poles = RIAA[0], Zeros = RIAA[1] and Poles,Zeros = RIAA
# both work. 

RIAA = [[50.048724,2122.0659],[500.48724,inf]]

def dB(x):
    return 20*log10(abs(x))

def f_warp(f,fs):
    if f == 0:
        out = 0.
    elif f == inf:
        out = inf
    else:
        out = ((fs/pi)*tan(pi*abs(f)/fs)*(f/abs(f)))
    return out

def f_from_z(z,fs):
    out = []
    for x in z:
        if x == -1:
            out.append(inf)
        else:
            out.append((fs/pi)*(1-x)/(1+x))
    return out
     

def z_from_f(f,fs):
    out = []
    for x in f:
        if x == inf:
            out.append(-1.) 
        else:
            out.append(((fs/pi)-x)/((fs/pi)+x))
    return out

def Fs_at_f(Poles,Zeros,f,norm = 0):
    ans = 1.
    for z in Zeros:
        if z == inf:
            ans = ans*1.0
        else:
            ans = ans*(z+1j*f)
    for p in Poles:
        if p == inf:
            ans = ans*1.0
        else:
            ans = ans/(p+1j*f)
    if norm:
        ans = ans/max(abs(ans))
    return ans

def Fz_at_f(Poles,Zeros,f,fs,norm = 0):
    omega = 2*pi*f/fs
    ans = 1.
    for z in Zeros:
        ans = ans*(exp(omega*1j)-z_from_f([z],fs))
    for p in Poles:
        ans = ans/(exp(omega*1j)-z_from_f([p],fs))
    if norm:
        ans = ans/max(abs(ans))
    return ans

def z_coeff(Poles,Zeros,fs,g,fg,fo = 'none'):
    if fg == inf:
        fg = fs/2
    if fo == 'none':
        beta = 1.0
    else:
        beta = f_warp(fo,fs)/fo
    a = poly(z_from_f(beta*array(Poles),fs))
    b = poly(z_from_f(beta*array(Zeros),fs))
    gain = 10.**(g/20.)/abs(Fz_at_f(beta*array(Poles),beta*array(Zeros),fg,fs))
    
    return (a,b*gain)

def biquad_filter(xin,z_coeff):
    a = z_coeff[0]
    b = z_coeff[1]
    xout = zeros(len(xin))
    xout[0] = b[0]*xin[0]
    xout[1] = b[0]*xin[1] + b[1]*xin[0] - a[1]*xout[0]
    
    for j in range(2,len(xin)):
        xout[j] = b[0]*xin[j]+b[1]*xin[j-1]+b[2]*xin[j-2]-a[1]*xout[j-1]-a[2]*xout[j-2]

    return xout
    
def write_txt(filename,y):
    f = open(filename,'w')
    for x in range(0,len(y)):
        f.write('%(num)1.14e\n'% {"num" : y[x]})    
    f.close()

def read_txt(filename):
    y = []
    with open(filename,'r')as f:
        for line in f:
            y.append(float(line))
    return array(y)

def read_dat(filename):
    f = open(filename,'r')
    line = f.readline()
    samplerate = int(line.rsplit()[3])
    line = f.readline()
    channels = int(line.rsplit()[2])
    if channels == 2:
        left = []
        right = []
        for line in f:
            x = line.rsplit()
            left.append(float(x[1]))
            right.append(float(x[2]))
        f.close()       
        return channels,samplerate,array(left),array(right)
    else:
        left = []
        for line in f:
            x = line.rsplit()
            left.append(float(x[1]))
        f.close()       
        return channels,samplerate,array(left)

def write_dat(filename, sr, left, right='none'):
    f = open(filename,'w')
    f.writelines('; Sample Rate %d\n'%sr)
    period = 1/float(sr)
    
    if right == 'none':
        f.writelines('; Channels 1\n')
        norm = max(abs(left))
        if norm == 0:
            norm = 1
        for x in range(0 , len(left)):
            f.writelines('%1.9e %1.9e\n'%(x*period,left[x]/norm))
    else:
        f.writelines('; Channels 2\n')
        x = max(abs(left))
        norm = y = max(abs(right))
        if x > y:
            norm = x
        if norm == 0:
            norm = 1
        for x in range(0 , len(left)):
            f.writelines('%1.9e %1.9e %1.9e\n'%(x*period,left[x]/norm,right[x]/norm))
            
    f.close()

def plot_fft_log(fs, data, diff = 'none', start = 0,stop = 0, delay = 0):

    scale = float(fs)/len(data)
    if start < scale:
        start = scale
    first_bin = int(round(start/scale))
    if stop > fs/2 or stop == 0:
        stop = fs/2
    last_bin = int(round(stop/scale))
    
    f = linspace(0, scale*((len(data)/2)-1), len(data)/2)[first_bin:last_bin]

    phase = 0.
    if delay:
        phase = linspace(0,360.*delay,len(data),endpoint = False)
        
    if diff == 'none':
        FsdB = dB(fft.fft(data))[first_bin:last_bin]
        Fsph = (angle(fft.fft(data),deg=1) + phase)[first_bin:last_bin]
    else:
        FsdB = (dB(fft.fft(data))-dB(fft.fft(diff)))[first_bin:last_bin]
        Fsph = (angle(fft.fft(data),deg=1) + phase - angle(fft.fft(diff),deg=1))[first_bin:last_bin]
        
    fig = plt.figure(figsize=(15,7.5))
    fig.suptitle("Gain and Phase vs Frequency")
    
    ax1 = fig.add_subplot(121)
    ax1.plot(f, FsdB)
    ax1.grid(True, which = "both", linestyle = "-")
    ax1.set_xscale ("log")
    ax1.set_ylabel("Gain - dB")
    ax1.set_xlabel("Frequency - Hz")
    
    ax2 = fig.add_subplot(122)        
    ax2.plot(f, Fsph)
    ax2.grid(True, which = "both", linestyle = "-")
    ax2.set_xscale ("log")
    ax2.set_xlabel("Frequency - Hz")
    ax2.set_ylabel("Phase - degrees")

    plt.show()

def plot_fft_lin(fs, data,diff = 'none', start = 0, stop = 0, delay = 0):
    
    scale = float(fs)/len(data)
    if start < scale:
        start = scale
    first_bin = int(round(start/scale))
    if stop > fs/2 or stop == 0:
        stop = fs/2
    last_bin = int(round(stop/scale))

    f = linspace(0, scale*((len(data)/2)-1), len(data)/2)[first_bin:last_bin]
    
    phase = 0.
    if delay:
        phase = linspace(0,360.*delay,len(data),endpoint = False)

    if diff == 'none':
        FsdB = dB(fft.fft(data))[first_bin:last_bin]
        Fsph = (angle(fft.fft(data),deg=1) + phase)[first_bin:last_bin]
    else:
        FsdB = (dB(fft.fft(data))-dB(fft.fft(diff)))[first_bin:last_bin]
        Fsph = (angle(fft.fft(data),deg=1) + phase - angle(fft.fft(diff),deg=1))[first_bin:last_bin]

    FsdB = dB(fft.fft(data))[first_bin:last_bin]
    Fsph = (angle(fft.fft(data),deg=1)+phase)[first_bin:last_bin]
                                        
    fig = plt.figure(figsize=(15,7.5))
    fig.suptitle("Gain and Phase vs Frequency")
    
    ax1 = fig.add_subplot(121)
    ax1.plot(f, FsdB)
    ax1.grid(True, which = "both", linestyle = "-")
    ax1.set_ylabel("Gain - dB")
    ax1.set_xlabel("Frequency - Hz")
    
    ax2 = fig.add_subplot(122)        
    ax2.plot(f, Fsph)
    ax2.grid(True, which = "both", linestyle = "-")
    ax2.set_xlabel("Frequency - Hz")
    ax2.set_ylabel("Phase - degrees")

    plt.show()

def plot_dat(sound):
    
    if len(sound) == 3:
        [channels, sample_rate, left] = sound

    if len(sound) == 4:
        [channels, sample_rate, left, right] = sound
    
    if channels == 1:
        time = float((len(left)-1))/sample_rate
        f = linspace(0, time , len(left))
        
        fig = plt.figure(figsize=(10,5))
        ax = fig.add_subplot(111)
        ax.plot(f, left)
        ax.grid(True, which = "both", linestyle = "-")
        ax.set_ylabel("Amplitude")
        ax.set_xlabel("Time, sec")
        plt.show()

    if channels == 2:
        time = float((len(left)-1))/sample_rate
        f = linspace(0, time , len(left))
        m1 = max(abs(left))
        maxv = m2 = max(abs(right))
        if m1 > m2:
            maxv = m1
        
        fig = plt.figure(figsize=(10,10))
        
        ax = fig.add_subplot(211)
        ax.axis([0,time,-maxv,maxv])
        ax.plot(f, left)
        ax.grid(True, which = "both", linestyle = "-")
        ax.set_ylabel("Amplitude")
        ax.set_title("Left Channel")
        ax.set_xlabel("Time, sec")
        
        ax = fig.add_subplot(212)
        ax.axis([0,time,-maxv,maxv])
        ax.plot(f, right)
        ax.grid(True, which = "both", linestyle = "-")
        ax.set_ylabel("Amplitude")
        ax.set_title("Right Channel")
        ax.set_xlabel("Time, sec")
        plt.show()

def plot_gain_phase(Poles, Zeros, zPoles = [], zZeros = [], fs = 48000.,fo = 1000. ,start = 20, stop = 20000,style = 0):

    flog = linspace(log10(start), log10(stop), 1000) 
    f = 10.0 ** flog

    Fs = dB(Fs_at_f(Poles,Zeros,f))-dB(Fs_at_f(Poles,Zeros,fo))
    angFs = angle(Fs_at_f(Poles,Zeros,f),deg = 1)
    if zPoles and zZeros:
        Fz = dB(Fz_at_f(Poles,Zeros,f,fs))-dB(Fz_at_f(Poles,Zeros,fo,fs))
        angFz = angle(Fz_at_f(Poles,Zeros,f,fs),deg = 1)

    fig = plt.figure(figsize=(10,10))
    ax1 = fig.add_subplot(211)
    if zPoles and zZeros:
        if style:
            ax1.plot(f, Fs-Fz)
        else:
            ax1.plot(f, Fs)
            ax1.plot(f, Fz)
    else:
        ax1.plot(f, Fs)
        

    ax1.grid(True, which = "both", linestyle = "-")
    ax1.set_xscale ("log")
    ax1.set_ylabel("Gain - dB")
    ax1.set_xlabel("Frequency - Hz")
    ax1.set_title("Gain and Phase Plot")

    ax2 = fig.add_subplot(212)
    if zPoles and zZeros:
        if style:
            ax2.plot(f, angFs-angFz)
        else:
            ax2.plot(f, angFs)
            ax2.plot(f, angFz)
    else:
        ax2.plot(f, angFs)
        
    ax2.grid(True, which = "both", linestyle = "-")
    ax2.set_xscale ("log")
    ax2.set_xlabel("frequency - Hz")
    ax2.set_ylabel("Degrees")

    plt.show()

def min_phase_FIR(Poles,Zeros,fs,outfile,plot = 0):
    
    #Initialize a complex array to hold the frequency domain response

    Fdomain = zeros(fs,dtype=complex)

    #The "trick". The phase must be rotated at each frequency to eliminate
    #the abrupt discontinuity between fs/2 and -fs/2 in the frequency domain.
    #This dramatically reduces ringing in the time domain substituting a slight
    #delay which does not hurt the minimum phase properties.

    phz = -angle(Fs_at_f(Poles,Zeros,fs/2))
    
    #Fill the array with the response up to fs/2 and fill the negative frequencies
    #with the complex conjugate which guarantees a real only inverse FFT. The phase
    #rotation is simply a multiply by e to the jTheta.

    for k in range(0,int(fs)/2-1):
        Fdomain[k+1]= exp(2j*phz*(float((k+1))/fs))*Fs_at_f(Poles,Zeros,k+1)
        Fdomain[fs-1-k] = conj(Fdomain[k+1])

    #Fs/2 and zero frequency are done separately as book keeping
    #The phase rotation always makes fs/2 have 0 phase.
    
    Fdomain[fs/2] = abs(Fs_at_f(Poles,Zeros,fs/2))
    Fdomain[0] = Fs_at_f(Poles,Zeros,0)

    #Now the impulse response is simply the inverse FFT. It is worth noting that
    #it is only very slightly non-causal in the case of RIAA.

    imp_res = fft.ifft(Fdomain)
    imp_res = array(real(imp_res),dtype = float64)

    #The result is normalized to 1, if you want you can normalize to 2*23-1 (24 bits).
    #I have not explored the implications with all programs, SoX does not compute a gain
    #normalization, but Foobar appears to fix the gain automatically so you don't clip.

    scale = max(abs(imp_res))
    imp_res = (imp_res/scale)

    #move the peak to the middle
    
    imp_res = roll(imp_res,int(fs/2))

    #Find where the output becomes less that 1/2 LSB

    mmax = 0
    mmin = 0
    for x in range(0,int(fs)):
        if (abs(imp_res[x]) > 2**-24) and (mmin == 0):
            mmin = x

    for x in range(int(fs)-1,0,-1):
        if (abs(imp_res[x]) > 2**-24) and (mmax == 0):
            mmax = x

    #write the trimmed file to disk
            
    write_dat(HOME + outfile + '.dat',fs,imp_res[mmin:mmax+1])

    if plot:
        fig = plt.figure(figsize=(10,5))
        ax = fig.add_subplot(111)

        f = linspace(0, mmax-mmin, mmax-mmin+1)
        ax.plot(f, imp_res[mmin:mmax+1])
        ax.grid(True, which = "both", linestyle = "-")
        ax.set_ylabel("Value")
        ax.set_title("Value vs. sample")
        plt.show()