"""
A skeleton for plotting fourier domain representation of signals
"""
import numpy as np
from numpy.fft import fft, fftshift, fftfreq
import matplotlib.pyplot as plt


def plot_spectrum(x, Ts, plot_x=False, in_f=True, xtick_count=16):
    """
    PLots magnitude spectrum of signal x with sampling period Ts
    Remember: omega = 2 pi f
    discrete omega = 2 pi f * T_s in [0, pi]
    """
    N = x.shape[0]

    if plot_x:
        plt.figure()
        plt.plot(np.arange(N), x)
        plt.xlabel('Discrete index n')
        plt.ylabel('Time domain x[n]')
        plt.grid()
        plt.show()

    # e.g. [-2, -1, 0, 1] Hz for N = 4 and Ts = 1/4
    sampled_frequencies = fftshift(fftfreq(N, Ts))
    # Sampled frequencies in discrete omega, mind the missing pi
    sampled_frequencies_d_omega = 2*sampled_frequencies*Ts

    xticks1 = np.linspace(np.min(sampled_frequencies),
                          np.max(sampled_frequencies), xtick_count)
    xticks2 = np.linspace(np.min(sampled_frequencies_d_omega),
                          np.max(sampled_frequencies_d_omega), xtick_count)
    xticks2 = list(map(lambda x: (str(round(x, 2)) + r'$\pi$'), xticks2))

    spectrum = fftshift(fft(x))

    plt.figure()
    plt.plot(sampled_frequencies, 10*np.log10(np.abs(spectrum)))
    plt.xlabel(r'Frequency $f$ in Hz')
    if in_f:
        plt.xticks(xticks1, rotation=30)
    else:
        plt.xticks(xticks1, xticks2, rotation=20)
    plt.ylabel(r'Magnitude of spectrum $10\log_{10}\|X[f]\|$ in dB')
    plt.grid()
    plt.show()


if __name__ == '__main__':
    N = 256
    Ts = 1e-3
    fs = 1/Ts
    nyquist_rate = fs/2
    t = np.arange(N)*Ts
    f = [50, 150, 250, 300]
    a = [10, 20, 30, 40]

    x = np.zeros(N)
    for i, fi in enumerate(f):
        assert fi <= nyquist_rate
        x += a[i]*np.cos(2*np.pi*fi*t)

    plot_spectrum(x, Ts, True, False)
