import numpy as np
from numpy import sin, cos, tan, pi
import csv
import sys
from time import time


x0 = 0.0
x1 = 1.0
ndata = 1024

outfile = 'dft.csv'


argv = sys.argv
narg = len(argv)
if narg >= 2:
    ndata = int(argv[1])


def usage():
    print("")
    print("Usage: python {} ndata".format(argv[0]))
    print("   ex: python {} {}".format(argv[0], ndata))
    print("")

def terminate():
    usage()
    exit()
    


# three trigonometric functions with randomized frequencies
# A: amplitude, f: frequency, p1: phase    
def func(x):
    x += 0.03 * np.random.rand();
    f1, p1, A1 = ( 1.5, pi/4.0, 1.0);
    f2, p2, A2 = ( 3.0, pi/3.0, 0.3);
    f3, p3, A3 = (10.0, pi/6.0, 0.5);
    return A1 * sin(2.0*pi * f1 * x + p1) \
         + A2 * sin(2.0*pi * f2 * x + p2) \
         + A3 * sin(2.0*pi * f3 * x + p3)

def fft(y):
    return np.fft.fft(y)

def dft(n, x, y, mode):
    if mode == 'inverse':
        C = -1.0
    else:
        C = 1.0

    f = [0.0]*n
    yr = [0.0]*n
    yi = [0.0]*n
    wr = [0.0]*n
    wi = [0.0]*n

    df = 1.0 / (x[n-1] - x[0]);
    wr[0] = 1.0
    wi[0] = 0.0
    phi = C * 2.0*pi / n
    wr0 = cos(phi)
    wi0 = sin(phi)
    wr[1] = wr0
    wi[1] = wi0
    for k in range(2, n):
        wr[k] = wr[k-1] * wr0 - wi[k-1] * wi0
        wi[k] = wr[k-1] * wi0 + wi[k-1] * wr0

    for k in range(0, int(n/2)):
        f[k]  = df * (k+1)
        yr[k] = 0.0
        yi[k] = 0.0
        for i in range(0, n):
            ik = (i * k) % n
            wr0 = wr[ik];
            wi0 = wi[ik];
# x*exp(2pik) = (R + i I)(cos(phi) + i sin(phi))
#      Rcos(phi) - Isin(phi) + i(Rsin(phi) + Icos(phi))
            yr[k] += y[i].real * wr0 - y[i].imag * wi0;
            yi[k] += y[i].real * wi0 + y[i].imag * wr0;
    return int(n / 2), f, yr, yi

def dft2(n, x, y, mode):
    if mode == 'inverse':
        C = -1.0
    else:
        C = 1.0

    f = [0.0]*n
    yr = [0.0]*n
    yi = [0.0]*n
    df = 1.0 / (x[n-1] - x[0]);
    for k in range(0, int(n/2)):
        f[k]  = df * (k+1)
        yr[k] = 0.0
        yi[k] = 0.0
        for i in range(0, n):
            phi  = C * 2.0*pi / n * i*k;
            cosp = cos(phi);
            sinp = sin(phi);
# x*exp(2pik) = (R + i I)(cos(phi) + i sin(phi))
#      Rcos(phi) - Isin(phi) + i(Rsin(phi) + Icos(phi))
            yr[k] += y[i].real * cosp - y[i].imag * sinp;
            yi[k] += y[i].real * sinp + y[i].imag * cosp;
    return int(n / 2), f, yr, yi


def main():
    global x0, x1, ndata
    global csvfile
    global outfile

    x = np.zeros(ndata)
    y = np.zeros(ndata)
    print("Make data by func(x) with random scattering")
    xstep = (x1 - x0) / ndata
    print("x = ({}, {}, {})".format(x0, x1, xstep))
    print("ndata={}".format(ndata))
    for i in range(0, ndata):
        x[i] = x0 + i * xstep
        y[i] = func(x[i])
    ndata = len(x)

    starttime = time()
    nDFT, fDFT, DFTr, DFTi  = dft(ndata, x, y, 1)
    endtime = time()
    print("Perform DFT (using rotation factor, faster): {} sec".format(endtime - starttime))
    print("")

    starttime = time()
    nDFT2, fDFT2, DFT2r, DFT2i = dft2(ndata, x, y, 1)
    endtime = time()
    print("Perform DFT2 (calculate sin/cos everytime, slow): {} sec".format(endtime - starttime))
    print("")

    nFFTiter = 10000
    starttime = time()
    for i in range(nFFTiter):
        FFT  = fft(y)
    endtime = time()
    print("Perform FFT(fast): {} sec".format((endtime - starttime) / nFFTiter))
    print("")

    print("Save FFTed data to [{}]".format(outfile))
    f = open(outfile, 'w')
    fout = csv.writer(f, lineterminator='\n')
    fout.writerow(['x', 'y', 'FT(DFT).r', 'FT(DFT).i',
                   'FT(DFT2).r', 'FT(DFT2).i', 'FT(FFT).r', 'FT(FFT).i'])
    for i in range(0, int(ndata/2)):
        fout.writerow([x[i], y[i], DFTr[i], DFTi[i], DFT2r[i], DFT2i[i], FFT[i].real, FFT[i].imag])


    terminate()
    

if __name__ == "__main__":
    main()
