import numpy as np
import matplotlib.pyplot as plt
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 add_noise_db(s, noise_db):
    var_noise = 10**(noise_db/10)
    noise = var_noise**0.5 * np.random.randn(len(s))
    return s + noise


def lowpass(data, f_cutoff, f_s):
    nyq = 0.5*f_s
    normal_cutoff = f_cutoff/nyq
    b, a = sp.signal.butter(1, normal_cutoff, btype="low", analog=False)
    y = sp.signal.lfilter(b, a, data)
    return y

def dsb_am(m, f_c, t):
    omega_c = 2*np.pi*f_c
    return (1 + m)*np.cos(omega_c*t)

def dsb_sc(m, f_c, t):
    omega_c = 2*np.pi*f_c
    c = np.cos(omega_c*t)
    return m*c

def product_demod(s, f_baseband, f_c, f_s, t):
    omega_c = 2*np.pi*f_c
    mixed = s*np.cos(omega_c*t)
    return lowpass(mixed, f_baseband, f_s)

def main():
    T_s = 0.0005
    f_s = 1/T_s

    f_m = 10
    omega_m = 2*np.pi*f_m

    f_c = 100

    t = np.arange(0,1,T_s)
    f = np.linspace(0,f_s,len(t))

    m = 0.5*(np.sin(omega_m*t) + np.sin(omega_m/2*t))
    s_am = dsb_am(m, f_c, t)
    s_sc = dsb_sc(m, f_c, t)

    m_am_demod = product_demod(s_am, f_m, f_c, f_s, t)
    m_sc_demod = product_demod(s_sc, f_m, f_c, f_s, t)

    plt.plot(t, m, label="Original Message Signal")
    plt.plot(t, m_am_demod, label="Demodulated DSB-AM Signal")
    plt.plot(t, m_sc_demod, label="Demodulated DSB-SC Signal")
    plt.legend(loc="upper right")
    plt.savefig("demodulation-clean.png")
    plt.show()

    m_am_noisy_demod = product_demod(add_noise_snr_db(s_am, 3), f_m, f_c, f_s, t)
    m_sc_noisy_demod = product_demod(add_noise_snr_db(s_sc, 3), f_m, f_c, f_s, t)

    plt.plot(t, m, label="Original Message Signal")
    plt.plot(t, m_am_noisy_demod, label="Demodulated DSB-AM Signal With Noise")
    plt.plot(t, m_sc_noisy_demod, label="Demodulated DSB-SC Signal With Noise")
    plt.legend(loc="upper right")
    plt.savefig("demodulation-noisy.png")
    plt.show()

if __name__ == "__main__":
    main()