import numpy as np
import matplotlib.pyplot as plt


def func(t):
    return -np.cos(2*np.pi*t)*(1+5*t**2)


# Numbers of sampling points
N = 4       # coarse mesh
Nf = 8      # fine mesh

tmax = 1.0
tfmax = 1.0

tstep = tmax / N
fstep = 1.0 / tmax
fmax = 1.0 / tstep

tfstep = tmax / Nf
ffstep = 1.0 / tfmax
ffmax = 1.0 / tfstep

t = np.arange(0, tmax, tstep)
# calculate frequencies: can be replaced with np.fft.fftfreq(N, d=t_step)
freq1 = np.arange(0, fmax, fstep)
# centering frequencies: can be replaced with np.fft.fftshift(freq)
freq2 = np.arange(-fmax / 2.0, fmax / 2.0, fstep)

tf = np.arange(0, tfmax, tfstep)
# calculate frequencies: can be replaced with np.fft.fftfreq(N, d=t_step)
freqf1 = np.arange(0, ffmax, ffstep)
# centering frequencies: can be replaced with np.fft.fftshift(freq)
freqf2 = np.arange(-ffmax / 2.0, ffmax / 2.0, ffstep)


f = func(t)
F = np.fft.fft(f)
Amp1 = np.abs(F)
# centering FTed data: can be replaced with np.fft.fftshift(F)
Amp2 = Amp1.tolist()[N//2:N].copy()
Amp2.extend(Amp1.tolist()[:N//2])

ff = func(tf)
Ff = np.fft.fft(ff)
Ampf1 = np.abs(Ff)
# centering FTed data: can be replaced with np.fft.fftshift(Ff)
Ampf2 = Ampf1.tolist()[Nf//2:Nf].copy()
Ampf2.extend(Ampf1.tolist()[:Nf//2])


# plot
fontsize = 16

fig, axes = plt.subplots(3, 1, figsize=(10, 8))
plt.rcParams["font.size"] = fontsize

ax = axes[0]
ax.plot(tf, ff, marker = "x")
ax.plot(t, f, linestyle = '', marker = "o")
ax.set_xlabel("Time (s)", fontsize = fontsize)
ax.set_ylabel("Signal", fontsize = fontsize)
ax.set_xlim(0, tmax)
#ax.grid()
ylim = ax.get_ylim()
ax.vlines(t,  *ylim, color = "red",  linestyles = "dashed", linewidth = 1.0)
ax.vlines(tf, *ylim, color = "blue", linestyles = "dotted", linewidth = 0.6)

ax = axes[1]
ax.plot(freqf1, Ampf1, marker = "x")
ax.plot(freq1, Amp1, marker = "o")
ax.set_xlabel("Frequency (Hz)", fontsize = fontsize)
ax.set_ylabel("Amplitude", fontsize = fontsize)
ax.set_xlim(0, ffmax)
#ax.grid()
ylim = ax.get_ylim()
ax.vlines(freq1,  *ylim, color = "red",  linestyles = "dashed", linewidth = 1.0)
ax.vlines(freqf1, *ylim, color = "blue", linestyles = "dotted", linewidth = 0.6)

ax = axes[2]
ax.plot(freqf2, Ampf2, marker = "x")
ax.plot(freq2, Amp2, marker = "o")
ax.set_xlabel("Frequency (Hz)", fontsize = fontsize)
ax.set_ylabel("Amplitude", fontsize = fontsize)
ax.set_xlim(-ffmax / 2.0, ffmax / 2.0)
#ax.grid()
ylim = ax.get_ylim()
ax.vlines(freq2,  *ylim, color = "red",  linestyles = "dashed", linewidth = 1.0)
ax.vlines(freqf2, *ylim, color = "blue", linestyles = "dotted", linewidth = 0.6)

plt.tight_layout()
plt.pause(1.0e-4)
input(">>")
