import numpy as np
from ERPSpace import ERBSpace
import pylab

def fft2gammatonemx(nfft, sr=16000, nfilts=64, width=1.0, minfreq=100, *othermore):
    sr=float(sr)
    nfilts=float(nfilts)
    minfreq=float(minfreq)
    EarQ = 9.26449
    minBW = 24.7
    order = 1.0
    
    if len(othermore)==0:
        maxfreq=sr/2
        maxlen=nfft
    elif len(othermore)==1:
        maxfreq=float(othermore[0])
        maxlen=nfft
    elif len(othermore)==2:
        maxfreq=float(othermore[0])
        maxlen=othermore[1]
    else:
        print "param error"
        return 0
    
    wts = np.zeros((nfilts, nfft),dtype=np.float64)
    cfreqs = ERBSpace(minfreq, maxfreq, nfilts)
    cfreqs = np.flipud(cfreqs)
    GTord = 4
    ucirc = np.exp(2j*np.pi*np.arange(nfft/2+1)/nfft)
    justpoles = 0

    for i in range(int(nfilts)):
        cf = cfreqs[i];
        ERB = width*((cf/EarQ)**order + minBW**order)**(1/order);
        B = 1.019*2*np.pi*ERB
        r = np.exp(-B/sr)
        theta = 2*np.pi*cf/sr
        pole = r*np.exp(1j*theta)

        if justpoles == 1:
            cosomegamax = (1.0+r*r)/(2.0*r)*np.cos(theta);
            if np.abs(cosomegamax) > 1:
                if theta < np.pi/2.0:
                    omegamax = 0 
                else:
                    omegamax =np.pi
            else:
                omegamax = np.arccos(cosomegamax)

            center = np.exp(1j*omegamax)
            gain = np.abs((pole-center)*(pole.conjugate()-center))**GTord
            wts[i,1:(nfft/2+1)] = gain * (np.abs((pole-ucirc)*(pole.conjugate()-ucirc))**(-GTord))
        else:
            T = 1/sr;
            A11 = -(2.0*T*np.cos(2.0*cf*np.pi*T)/np.exp(B*T) + 2.0*np.sqrt(3.0+2.0**1.5)*T*np.sin(2*cf*np.pi*T)/np.exp(B*T))/2.0 
            A12 = -(2.0*T*np.cos(2.0*cf*np.pi*T)/np.exp(B*T) - 2.0*np.sqrt(3.0+2.0**1.5)*T*np.sin(2*cf*np.pi*T)/np.exp(B*T))/2.0
            A13 = -(2.0*T*np.cos(2.0*cf*np.pi*T)/np.exp(B*T) + 2.0*np.sqrt(3.0-2.0**1.5)*T*np.sin(2*cf*np.pi*T)/np.exp(B*T))/2.0 
            A14 = -(2.0*T*np.cos(2.0*cf*np.pi*T)/np.exp(B*T) - 2.0*np.sqrt(3.0-2.0**1.5)*T*np.sin(2*cf*np.pi*T)/np.exp(B*T))/2.0 
            zros = -np.array([A11,A12,A13,A14])/T
    
            gain =  np.abs((-2.0*np.exp(4j*cf*np.pi*T)*T + 2.0*np.exp(-(B*T) + 2j*cf*np.pi*T)*T* \
                             (np.cos(2.0*cf*np.pi*T) - np.sqrt(3.0 - 2.0**(3.0/2.0))* np.sin(2.0*cf*np.pi*T)))* \
                             (-2.0*np.exp(4j*cf*np.pi*T)*T + 2.0*np.exp(-(B*T) + 2j*cf*np.pi*T)*T* \
                             (np.cos(2.0*cf*np.pi*T) + np.sqrt(3.0 - 2.0**(3.0/2.0))*np.sin(2.0*cf*np.pi*T)))* \
                             (-2.0*np.exp(4j*cf*np.pi*T)*T + 2.0*np.exp(-(B*T) + 2j*cf*np.pi*T)*T*(np.cos(2*cf*np.pi*T) - \
                             np.sqrt(3.0 + 2.0**(3.0/2.0))*np.sin(2.0*cf*np.pi*T)))*(-2.0*np.exp(4j*cf*np.pi*T)*T + \
                             2.0*np.exp(-(B*T) + 2j*cf*np.pi*T)*T*(np.cos(2*cf*np.pi*T) + np.sqrt(3.0 + 2.0**(3.0/2.0))* \
                             np.sin(2.0*cf*np.pi*T)))/(-2.0/np.exp(2*B*T) - 2.0*np.exp(4j*cf*np.pi*T)+2.0*(1 + np.exp(4j*cf*np.pi*T))/np.exp(B*T))**4)
            
            wts[i,np.arange(0,int(nfft/2+1))] = ((T**4)/gain)*np.abs(ucirc-zros[0])*np.abs(ucirc-zros[1])*np.abs(ucirc-zros[2])*np.abs(ucirc-zros[3]) \
                                                *(np.abs((pole-ucirc)*(pole.conjugate()-ucirc))**(-GTord))

    wts = wts[:,0:maxlen]
    return wts

# Testing
if __name__ == '__main__':
   
    cfs = fft2gammatonemx(256, 16000, 64,1,50,8000, 1024.0/2.0+1);
    print cfs
    pylab.figure()
    pylab.imshow(cfs, origin='lower', aspect='auto', interpolation='nearest')
    pylab.ylabel('Frequency')
    pylab.show()