#! /usr/bin/python
##! /Users/dhw/anaconda/bin/python
# fftsimple.py -- generate a sinusoidal function and measure/plot its FFT
# fftsimple.py n freq x0
#   n = length of array
#   freq = frequency of sine wave in cycles per unit x
#   x0 = starting point of sine wave
# function is y = sin(2*pi*freq*x+offset) where x runs from 0 to 1

import sys
import numpy as np
import matplotlib.pyplot as plt
from math import *

TOL=1.e-4

n=int(sys.argv[1])
freq=float(sys.argv[2])
x0=float(sys.argv[3])

x=np.arange(n)/float(n)
y=np.cos(2*pi*freq*x+x0)
xfine=np.arange(1000)/1000.
yfine=np.cos(2*pi*freq*xfine+x0)

# Note division by n here, so what's eventually plotted is FT/n
z=np.fft.fft(y)/n
#z=np.fft.rfft(y)/n		# uncomment this for real-to-complex fft
f=np.arange(z.size)
zr=np.real(z)
zi=np.imag(z)

# print all Fourier modes whose modulus is larger than the value TOL
modz=np.absolute(z)
out=np.array([f[modz>TOL],zr[modz>TOL],zi[modz>TOL],modz[modz>TOL]])
np.set_printoptions(precision=3)
print out.transpose()

plt.subplot(2,1,1)
plt.plot(x,y)
plt.plot(x,y,'bo')
plt.plot(xfine,yfine,'r')

plt.subplot(2,1,2)
plt.plot(f,zr,'ro',f,zi,'b^')

plt.show()
#plt.savefig('simple.eps')