import numpy as np
import matplotlib.pyplot as plt
import scipy as sp


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
    c = np.cos(omega_c*t)
    return 0.5*m*c + c

def dsb_sc(m, f_c, t):
    omega_c = 2*np.pi*f_c
    c = np.cos(omega_c*t)
    return 0.5*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 = 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.show()

if __name__ == "__main__":
    main()