#!/usr/bin/python3
import sys
import numpy as np
import matplotlib.pyplot as plt
import wave
import struct
import scipy.io.wavfile
from scipy import signal

# generate original square wave
def square_wave(T0=10, Tpulse=2/3, A=1, dur=0.3, fs=44100):
    t = np.linspace(0, dur, fs*dur, False)
    swave = (A/2)*signal.square((2*np.pi*T0*(t+0.1/3)), duty = Tpulse)+0.5
    return swave

# calculate amplitudes for spectral signal
def amplitudes(n):
    a = []
    a.append(1/3)
    for i in range (1,n+1):
        a.append( (np.sin((2*np.pi*i)/3)) / (np.pi*i) )
    return a

# calculate frequencies of harmonics
def frequencies(n, f0=10):
    f = []
    for i in range(0,n+1):
        f.append(i*f0)
    return f

# generate periodic sound
def periodic(t, f, A, phi):
    x = A*(np.cos((2*np.pi*f*t)+phi*np.pi))
    return x

# plot two waves on the same graph
def plot_two_waves(y1, y2, n, filename, fs=44100):
    t = np.linspace(0, len(y1)/fs, len(y1), False)
    plt.plot(t, y1)
    plt.plot(t, y2)
    plt.title('Square wave')
    plt.xlabel('Time(s)')
    plt.ylabel('X')
    plt.grid(True)
    plt.ylim(-0.2,1.2)
    plt.savefig(filename)
    plt.show()

# plot the spectrum using a lollipop graph
def plot_spectrum(n):
    spectrum = []
    t = np.linspace(-n, n, (n*2)+1, False)
    for i in range (-n,0):
        spectrum.append( (np.sin((2*np.pi*i)/3)) / (np.pi*i) )
    spectrum.append(2/3)
    for i in range (1,n+1):
        spectrum.append( (np.sin((2*np.pi*i)/3)) / (np.pi*i) )
    plt.title('Spectrum')
    plt.xlabel('N')
    plt.ylabel('Xn')
    plt.stem(t,spectrum,use_line_collection=True)
    plt.show()

# generate the reconstructed wave
def generate_wave(amps, freqs, phi=0, dur=0.3, fs=44100):
    sum = 0
    t = np.linspace(0, dur, fs*dur, False)
    for x in range (0, len(amps)):
        z = 2*periodic(t, freqs[x], amps[x], phi)
        sum += z
    return sum

# normalise to 16 bits
def normalize(y):
    for x in range(0,len(y)):
        y[x] = y[x]/3.2767
        if (y[x] > 1):
            y[x] = 1
    y = np.array(y, dtype='int')

def write_sound(y, filename="sound", fs=44100, nchannels=1):
    scipy.io.wavfile.write('filename.wav',fs,y)

if __name__ == '__main__':

    # get n
    n = int(sys.argv[1])

    # get file name
    filename = sys.argv[2]

    # create array of amplitudes for spectral components
    amps = amplitudes(n)

    # create array of frequencies for spectral components
    freqs = frequencies(n)

    # generate reconstructed wave that is the sum of all spectral components
    r_wave = generate_wave(amps, freqs)

    # generate original square wave
    s_wave = square_wave()

    # plot original square wave and reconstructed wave on same graph
    plot_two_waves(r_wave, s_wave, n, filename)

    # plot spectrum on separate graph
    plot_spectrum(n)

    # normalize the square wave
    normalize(s_wave)

    # write the square wave to a sound file
    write_sound(s_wave)