import matplotlib
matplotlib.use( 'TKAgg' )

import time

from numpy import fft
import numpy as np
import matplotlib.pyplot as plt

# rosettacode FFT - strange workings.. 

def rosetta_fft(x):
    """A recursive implementation of the 1D Cooley-Tukey FFT"""
    x = np.asarray(x, dtype=float)
    N = x.shape[0]
    
    if N % 2 > 0:
        raise ValueError("size of x must be a power of 2")
    elif N <= 32:  # this cutoff should be optimized
        return DFT_slow(x)
    else:
        X_even = rosetta_fft(x[::2])
        X_odd = rosetta_fft(x[1::2])
        factor = np.exp(-2j * np.pi * np.arange(N) / N)
        return np.concatenate([X_even + factor[:N / 2] * X_odd,
                               X_even + factor[N / 2:] * X_odd])
                               

def DFT_slow(x):
    """Compute the discrete Fourier Transform of the 1D array x"""
    x = np.asarray(x, dtype=float)
    N = x.shape[0]
    n = np.arange(N)
    k = n.reshape((N, 1))
    M = np.exp(-2j * np.pi * k * n / N)
    return np.dot(M, x)
    
plt.ion()

#sampling parameters
N=100
xmin=0.
xmax=20.
dx=(xmax-xmin)/N
print "sampling freq:",1./dx
nuc=0.5/dx
print "critical freq:",nuc
#x=np.linspace(xmin,xmax,N)
x=np.linspace(xmin,xmax,N,endpoint=False)
#print x

#signal parameters 
lambda1=10.
lambda2=5.
k1=1./lambda1
k2=1/lambda2

#case when there is aliasing in the sampling.. nuc=0.5
#k1=0.30
#k2=0.70

print "freqs",k1,k2
fx=np.sin(2*np.pi*x*k1)+np.sin(2*np.pi*x*k2)
#fx=np.sin(2*np.pi*x*k2)


plt.cla()
plt.plot(x,fx,'r-o')
raw_input( "Press Enter to continue... " )

plt.cla()
Fk = fft.fft(fx)/N # Fourier coefficients (divided by n) 
ffx=fft.ifft(Fk*N) # immediate inverse

#Fk_rosetta=rosetta_fft(fx)
#print Fk
#print Fk_rosetta/N

nu = fft.fftfreq(N,dx) # Natural frequencies
Fk = fft.fftshift(Fk) # Shift zero freq to center
nu = fft.fftshift(nu) # Shift zero freq to center
f, ax = plt.subplots(3,1,sharex=True)
# Plot Cosine terms
ax[0].plot(nu, np.real(Fk),'r-o') 
ax[0].set_ylabel(r'$Re[F_k]$', size = 'x-large') 
# Plot Sine terms
ax[1].plot(nu, np.imag(Fk),'g-o') 
ax[1].set_ylabel(r'$Im[F_k]$', size = 'x-large')
# Plot spectral power
ax[2].plot(nu, np.absolute(Fk)**2,'-o') 
ax[2].set_ylabel(r'$\vert F_k \vert ^2$', size = 'x-large')
ax[2].set_xlabel(r'$\widetilde{\nu}$', size = 'x-large') 

raw_input( "Press Enter to continue... " )


f, ax = plt.subplots(3,1,sharex=True)
# Plot Cosine terms
ax[0].plot(x, np.real(ffx),'r-o') 
ax[0].set_ylabel(r'$Re[ffx]$', size = 'x-large') 
# Plot Sine terms
ax[1].plot(x, np.imag(ffx),'g-o') 
ax[1].set_ylabel(r'$Im[ffx]$', size = 'x-large')
# Plot spectral power
ax[2].plot(x, np.real(ffx)-fx,'-o') 
ax[2].set_ylabel(r'$Re[ffx]-fx$', size = 'x-large')
ax[2].set_xlabel(r'$x$', size = 'x-large') 


raw_input( "Press Enter to continue... " )