#!/usr/bin/env python

"""Demonstrate various aspects of the FFT as implemented in numpy and scipy"""


from numpy import *
from numpy.fft import fft, fftfreq, ifft
import Gnuplot, Gnuplot.funcutils
import time

k = 3.    #wave number
N = 1024   #samples
L = 10.   #lenght of interval
signal = zeros(N)
#signal[:N/10] = 1
#signal[2*N/10:3*N/10] = 1
signal[4*N/10:5*N/10] = 1
signal[5*N/10:6*N/10] = 1
#signal[6*N/10:7*N/10] = 1
x = linspace(0.,L,N)
#ran = random.standard_normal(N)
#signal = 3.*sin(k*2.*pi/L*x)*(1.+0.1*ran) #+ 2.*sin((k+1)*2*pi/L*x)*(1.+0.1*random.standard_normal(N)) + 8.*sin((k+2)*2*pi/L*x)*(1.+0.03*random.standard_normal(N))
#signal = sin(k*2.*pi/L*x) +0.1*sin(100*2.*pi/L*x)
g = Gnuplot.Gnuplot()
#g('set data style line')
signal_data = Gnuplot.Data(x,signal,with_='line ')
g.plot(signal_data,title='signal')
raw_input('Press return to continue!')

t0 = time.clock()
fourier = fft(signal)
t1 = time.clock()
dt_np = t1-t0
print 'Numpy fft took ', dt_np, 'seconds'
timestep = L/N
freq = L*fftfreq(N, d=timestep)
#df = Gnuplot.Data(freq,with_='line')
#g.plot(df,title='frequencies')
#raw_input('Press return to continue!')
dat = Gnuplot.Data(freq,abs(fourier)**2,with_='line')
#g('set yrange [0.1:100000]')
g('set logscale y')
g.plot(dat,title='fft')
raw_input('Press return to continue!')

#inverse fft
filtered_signal = ifft(fourier)
fil_sig_dat = Gnuplot.Data(x,filtered_signal,with_='line')
g('unset logscale y')
g.plot(fil_sig_dat,signal_data,title='filtered and original signal')
raw_input('Press return to continue!')