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

Nsub = 100 # number of subsequences

w_1 = 20 # frequency of the 1st component of the signal (Hz)
w_2 = 40 # frequency of the 2nd component of the signal (Hz)

a = 1.1 # magnitude of the 1st component of the signal
b = 0.6 # magnitude of the 2nd component of the signal

t = np.array([i for i in range(1,301)])/1000 # time samples (s)
fs = 1 / (t[1]-t[0]) # sampling frequency (Hz)

x = a*np.cos(2*np.pi*w_1*t) + b*np.sin(2*np.pi*w_2*t) # considered signal

N = [len(x)-10, len(x), len(x)+10]
fig, ax = plt.subplots(len(N), 1, constrained_layout=True,\
                       figsize=(6, 7), dpi=250)

for idx, item in enumerate(N):
    FFT = np.fft.fft(x, n=item)
    amps = np.abs(FFT) / (len(FFT) / 2)
    ax[idx].stem(f[:20], amps[:20])
    ax[idx].grid(True)
    ax[idx].set_ylabel('Magnitude of the FFT')
plt.xlabel('Frequencies (Hz)')
plt.show()