import numpy as np
import sympy as sy
import matplotlib.pyplot as plt


t_sy = sy.Symbol("t", real=True)
fi_sy = 2 + 2*t_sy  # instantanenous frequency
phi_sy = fi_sy.integrate(t_sy) # instantaneous

fi_np = sy.lambdify(t_sy, fi_sy)  # numpy functions
phi_np = sy.lambdify(t_sy, phi_sy)
lb_fi = "$f_i(t)=%s$" % sy.latex(fi_sy)  # plot label

n_x = 2**10
T_x = 10/n_x
t_x = np.arange(n_x)*T_x

x = np.sin(2*np.pi*phi_np(t_x))  # generate signals
f_Xos, Xos = np.fft.rfftfreq(n_x*2, T_x), np.fft.rfft(x, n=2*n_x)

# Parameter for specgram:
nfft = 256 # Parameter NFFT
sargs = dict(x=2/nfft * x, NFFT=nfft, Fs=1/T_x, noverlap=nfft-1,
                   mode="magnitude", scale="linear", cmap=plt.cm.viridis)
fg1, ax1 = plt.subplots(1, 1)
fg2, ax2 = plt.subplots(1, 1)
ax1.set_title(r"256 Point Spectrogram with default 'xextent' parameter")
_, _, t_s1, im1 = ax1.specgram(**sargs)
ax2.set_title(r"256 Point Spectrogram with correct 'xextent' parameter")
_, _, t_s2, im2 = ax2.specgram(**sargs, xextent=(t_s1[0], t_s1[-1]))

for fg, ax, im in [(fg1, ax1, im1), (fg2, ax2, im2)]:
    ax.plot(t_x, fi_np(t_x), "r--", label=lb_fi, linewidth=.5)
    cb = fg.colorbar(im)
    cb.set_label("Magnitude $|X(t,f)|$")
    ax.legend()
    ax.set_xlim(t_x[0], t_x[-1])
    ax.set_xlabel(r"Time $t$ in Seconds ($\Delta t = %.3f\,$s)" % t_x[1])
    ax.set_ylabel(r"Frequency $f$ in Hertz ($\Delta f = %.3f\,$Hz)" % f_Xos[1])

plt.show()