#!/usr/bin/python3

import math

try:
        import matplotlib.pyplot as plt
        has_matplotlib = True
except:
        has_matplotlib = False

import argparse

defaultParameters = {
        'window_size': 1,
        'num_points': 1024,
}

class Complex:
        def __init__(self, real, im):
                self.real = real
                self.im   = im

        def __add__(self, n):
                if type(n) is Complex:
                        return Complex(self.real+n.real, self.im+n.im)
                else:
                        return Complex(self.real+n, self.im)

        def __mul__(self, n):
                if type(n) is Complex:
                        return Complex(self.real*n.real - self.im*n.im, self.real*n.im + self.im*n.real)
                else:
                        return Complex(self.real*n, self.im*n)

        def __truediv__(self, n):
                if type(n) is Complex:
                        m2 = n.real*n.real + n.im*n.im
                        return self*Complex(n.real/m2, -n.im/m2)
                else:
                        return Complex(self.real/n, self.im/n)

        def __str__(self):
                return "("+str(self.real)+","+str(self.im)+")"

        def mod(self):
                return math.sqrt(self.real*self.real + self.im*self.im)

        def e(v):
                return Complex(math.cos(v), math.sin(v))

def create_samples(window, n_samples, function):
        offset = window/n_samples
        samples = list()
        for i in range(0, n_samples):
                v = i*offset
                samples.append(function(v))
        return samples

def perform_dft(t_samples):
        n_samples = len(t_samples)
        f_samples = [Complex(0, 0)]*n_samples
        for k in range(0, n_samples):
                for n in range(0, n_samples):
                        v = -2*math.pi*k*n/n_samples
                        f_samples[k] += (Complex.e(v)*t_samples[n])
        return f_samples

sines = lambda x: 3+2*math.sin(50*x)+math.cos(90*x)
def ramp(x):
        if x < math.pi/2:
                return x*2/math.pi
        if x < 3*math.pi/2:
                return 2-x*2/math.pi
        if x < 2*math.pi:
                return -4+x*2/math.pi

def parseArgs():
        parser = argparse.ArgumentParser(description='Playing with DFT')
        parser.add_argument('-w', '--window', help='Window size (default = %.2f)' % defaultParameters['window_size'], default=defaultParameters['window_size'], type=float)
        parser.add_argument('-n', '--num_points', help='Number of points (default = %d)' % defaultParameters['num_points'], default=defaultParameters['num_points'], type=int)
        parser.add_argument('-p', '--plot', help='Plot time and frequency series', action='store_true')
        args = parser.parse_args()
        return args

def plot(args, t_samples, f_samples):
        # Drop all samples for frequencies higher than Fsampling/2
        m = [round(v.mod(), 3) for v in f_samples[0:int(args.num_points/2)]]

        f, axa = plt.subplots(2)

        # Get current axes
        axa[0].set_title('DFT example')

        axa[0].set_xlim([-0.5, args.window + 0.5])
        axa[0].set_xlabel('Time')
        axa[0].plot([x*args.window/args.num_points for x in list(range(0, args.num_points))], t_samples, color = 'b')
        axa[0].grid()

        delta_freq = 2*math.pi/args.window_size
        x = [v*delta_freq for v in list(range(0, int(args.num_points/2)))]
        axa[1].set_xlim([-0.5, x[-1]+0.5])
        axa[1].set_ylabel('Power Spectral Density')
        axa[1].set_xlabel('Frequency')
        axa[1].stem(x, m, 'r-')
        axa[1].grid()

        plt.show()

def show_dft(args, f_samples):
        delta_freq = 2*math.pi/args.window
        print("+---------+---------+")
        print("| %-7s | %-7s |" % ("Freq", "PSD"))
        print("+---------+---------+")

        for i, v in enumerate(f_samples[0:int(args.num_points/2)]):
                print("| %7.3f | %7.3f |" %(round(i*delta_freq,3), round(v.mod(),3)))

        print("+---------+---------+")

# Main program !
f=sines
args = parseArgs()
t_samples = create_samples(args.window, args.num_points, f)
f_samples = perform_dft(t_samples)

# Normalize value, then take only half the spectrum
for i,v in enumerate(f_samples):
        f_samples[i] = f_samples[i]*(2/args.num_points)
f_samples[0] = f_samples[0]/2

if args.plot:
        if has_matplotlib:
                plot(args, t_samples, f_samples)
        else:
                print("Plotting is not possible. Please install matplotlib.")
else:
        show_dft(args, f_samples)