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

f_m = 5E3
f_c = 25E3

f_s = 50*f_c

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)
f = t*f_s/(2*T_m)

# ===== AM =====
s = A_c * (1 + A_m*np.cos(2*np.pi*f_m*t)) * np.cos(2*np.pi*f_c*t)

var_s = np.cov(s)

SNR = 10
var_snr = var_s/(10**(SNR/10))
noise_snr = (var_snr**0.5) * np.random.randn(len(s))
m = s+noise_snr

S = sp.fft.fft(s)
M = sp.fft.fft(m)

plt.subplot(211)
plt.stem(f, S)
plt.subplot(212)
plt.stem(f, M)
plt.show()

# ===== FM =====
beta_f = 10
s = A_c*np.cos(2*np.pi*f_c*t + beta_f*A_m*np.sin(2*np.pi*f_m*t))

var_snr = var_s/(10**(SNR/10))
noise_snr = (var_snr**0.5) * np.random.randn(len(s))
m = s+noise_snr

S = sp.fft.fft(s)
M = sp.fft.fft(m)

plt.subplot(211)
plt.stem(f, S)
plt.subplot(212)
plt.stem(f, M)
plt.show()