import numpy as np
import matplotlib.pyplot as plt


def prep_data(freqs, phases, sample_freq, N):
    f1, f2 = freqs
    p1, p2 = phases
    t = 1 / sample_freq

    data = []
    for n in range(N):
        sin1 = np.sin(2 * np.pi * 1000 * n * t + p1)
        sin2 = 5 * np.sin(2 * np.pi * 2000 * n * t + p2)
        sin_d = sin1 + sin2
        data.append(sin_d)
    return data


def dft(data):
    N = len(data)
    dft_out = []
    for k in range(N):
        sigs = 0
        for n in range(N):
            sigs += data[n] * np.power(np.e, - 2 * np.pi * 1j * n * k / N)
        dft_out.append(sigs)
    return dft_out


def magnitude(x):
    return np.sqrt(np.power(x.real, 2) + np.power(x.imag, 2))


def phase(x):
    return np.arctan(x.imag / x.real)



N = 80
freqs = [1000, 2000]
phases = [0, 3 * np.pi / 4]
sample_freq = 8000

signal = prep_data(freqs, phases, sample_freq, N)
dft_out = dft(signal)
dft_mags = [magnitude(k) for k in dft_out]
dft_freqs = [phase(j) for j in dft_out]

fig, plots = plt.subplots(3, 1)

plots[0].set_title('Signal')
plots[0].plot(signal, color='red')

plots[1].set_title('DFT Frequency Magnitudes')
plots[1].plot(dft_mags, color='blue')

plots[2].set_title('DFT Frequency Phases')
plots[2].plot(dft_freqs, color='green')

plt.tight_layout()
plt.show()