import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from cycler import cycler
import scipy as sp


def add_noise(s, SNR):
    var_s = np.cov(s)
    var_noise = var_s/(10**(SNR/10))
    noise = var_noise**0.5 * np.random.randn(len(s))
    return s + noise

def plot_at_snrs(t, f, s_am, s_fm):
    stem_colors = ['k', 'b', 'r']
    ci = 0
    stem_color = stem_colors[ci]

    am_time = plt.figure()
    am_freq = plt.figure()
    fm_time = plt.figure()
    fm_freq = plt.figure()

    for SNR in range(10,30,10):
        m_am = add_noise(s_am,SNR)
        m_fm = add_noise(s_fm,SNR)

        plt.figure(am_time)
        plt.plot(t, m_am, label=f"SNR = {SNR}")

        plt.figure(am_freq)
        plt.stem(f, abs(sp.fft.fft(m_am)), label=f"SNR = {SNR}", markerfmt=stem_color+"o", linefmt=stem_color+"-")

        plt.figure(fm_time)
        plt.plot(t, m_fm, label=f"SNR = {SNR}")

        plt.figure(fm_freq)
        plt.stem(f, abs(sp.fft.fft(m_fm)), label=f"SNR = {SNR}", markerfmt=stem_color+"o", linefmt=stem_color+"-")

        ci += 1
        stem_color = stem_colors[ci]

    plt.figure(am_time)
    plt.plot(t, s_am, label="Pure signal")
    plt.legend(loc="upper right")
    plt.savefig("part-3-am-time")

    plt.figure(am_freq)
    plt.stem(f, abs(sp.fft.fft(s_am)), label="Pure signal", markerfmt=stem_color+"o", linefmt=stem_color+"-")
    plt.xlim((0, 2E5))
    plt.legend(loc="upper right")
    plt.savefig("part-3-am-freq")

    plt.figure(fm_time)
    plt.plot(t, s_fm, label="Pure signal")
    plt.legend(loc="upper right")
    plt.savefig("part-3-fm-time")

    plt.figure(fm_freq)
    plt.stem(f, abs(sp.fft.fft(s_fm)), label="Pure signal", markerfmt=stem_color+"o", linefmt=stem_color+"-")
    plt.xlim((0, 2E5))
    plt.legend(loc="upper right")
    plt.savefig("part-3-fm-freq")

    plt.show()


def main():
    f_m = 5E3
    f_c = 125E3
    f_s = 125E4

    T_m = 1/f_m
    T_c = 1/f_c
    T_s = 1/f_s

    A_c = 10
    A_m = 1

    t = np.arange(0,2*T_m,T_s)      # Time domain
    f = t*f_s/(2*T_m)               # Frequency domain

    beta_f = 10                     # Modulation index

    s_am = A_c * (1 + A_m*np.cos(2*np.pi*f_m*t)) * np.cos(2*np.pi*f_c*t)
    s_fm = A_c*np.cos(2*np.pi*f_c*t + beta_f*A_m*np.sin(2*np.pi*f_m*t))
    plot_at_snrs(t, f, s_am, s_fm)



if __name__ == "__main__":
    mpl.rcParams['axes.prop_cycle'] = cycler(color=['k', 'b', 'r'])
    main()