#!/usr/bin/env python3

import numpy as np

# Example Functions
def func1(x):
    return np.exp(20j * x) + np.exp(40j * x )

def func2(x):
    return np.exp( 5j * x**2 )

####################
# Fourier Transforms
####################
def DFT(yk):
    N = np.size(yk)
    xk = 2*np.pi/N*np.arange(N)

    beta = np.zeros(N, dtype=np.complex128)
    for j in range(N):
        beta[j] = np.dot(yk, np.exp(- 1j * j * xk ))

    return beta

def FFT(yk):
    N = np.size(yk)

    if ( N%2 > 0 ):
        raise ValueError("The data set does not consist of 2^M values:", N)

    if N <= 2:
        return DFT(yk)

    else:
        beta_even = FFT(yk[::2])
        beta_odd  = FFT(yk[1::2])

        exp_terms = np.exp( -2j * np.pi * np.arange(0,N) / N )

        beta_even_full = np.concatenate([beta_even, beta_even])
        beta_odd_full  = np.concatenate([beta_odd, beta_odd])

        return beta_even_full + exp_terms * beta_odd_full



# Plotting functions
def plot_DFTs(N = 128, fft1 = DFT, fft2 = np.fft.fft, datafuncs=[func1,func2]):
    from matplotlib import pyplot

    j = np.arange(N)
    xk = 2*np.pi/N * np.arange(N)
    pyplot.subplot()
    pyplot.title("Discrete Fourier Transform")
    pyplot.xlabel("j")
    pyplot.ylabel("|$\\beta_j$|")

    # First Function
    for i, dataf in enumerate(datafuncs):
        pyplot.plot(j, np.abs(fft1(func1(xk))), label="{} {}".format(fft1.__name__, dataf.__name__), linewidth=3)
        pyplot.plot(j, np.abs(fft2(func1(xk))), '--', label="{} {}".format(fft1.__name__, dataf.__name__))


    pyplot.legend()
    pyplot.show()


def time_FT_func(func, yk, *args, **kwargs):
    import timeit

    tOut = timeit.repeat(stmt=lambda: func(yk), *args, **kwargs)

    return np.mean(tOut)

def plot_FT_func_timing(func = DFT, M_max = 4, repeat = 10, number = 5, show = True):
    from matplotlib import pyplot

    M = np.arange(1,M_max)
    Ns = 2**M

    times = np.zeros(M_max-1)
    for i in range(np.size(times)):
        xk = 2*np.pi/Ns[i] * np.arange(Ns[i])
        yk = func1(xk)

        times[i] = time_FT_func(func, yk, repeat=repeat,number=number)

    pyplot.subplot()
    pyplot.title("Evaluation Time for {}*{} DFTs for each N ".format(repeat, number))
    pyplot.ylabel("Time (sec)")
    pyplot.xlabel("$N^2$")
    pyplot.ticklabel_format(scilimits=(0,0))

    pyplot.plot(Ns**2, times, "o:", label=func.__name__)
    pyplot.legend()

    if show:
        pyplot.show()



if __name__ == "__main__":
    plot_DFTs()
    plot_DFTs(fft1=DFT,fft2=FFT)

    plot_FT_func_timing(M_max = 9, func=FFT, show = False) # 12 takes 10 seconds, 13 takes 50 seconds
    plot_FT_func_timing(M_max = 9, func=DFT) # 12 takes 10 seconds, 13 takes 50 seconds