import sys
import numpy as np
from numpy import sin, cos, tan, sinh, cosh, tanh, exp, log, sqrt, pi
from pprint import pprint


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
