import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as sig

#%%## AM signal generation ###################################################
duration = 1 # in seconds
fs = 44100 # Hz
carrier_freq = 2000 # Hz
mod_freq = 200
t = np.linspace(0,duration,endpoint=True,num=fs)
audio = (1 + 1 * np.cos(2*np.pi*mod_freq*t)) * np.cos(2*np.pi*carrier_freq*t)

#%%###########################################################################
f, _, Zxx = sig.stft(audio, noverlap=(256 - 32), fs=fs)
aZxx = np.abs(Zxx)

#%%## Plot synchrosqueezed CWT ###############################################
plt.imshow(aZxx, aspect='auto', cmap='bone')
plt.show()
#%%## Find carrier freq ######################################################
max_row_idx = np.where(aZxx == aZxx.max())[0]
max_row = aZxx[max_row_idx].squeeze()

# print peak's frequency
print(f[max_row_idx])
#%%###########################################################################
# plot amplitude modulator
plt.plot(max_row[:800]); plt.show()

# find modulator frequency
max_row_fft = np.abs(np.fft.rfft(max_row))
plt.plot(max_row_fft); plt.show()

# peak at 200; exclude dc term
peak_bin = np.argmax(max_row_fft[1:]) + 1
print(peak_bin)