#!/usr/bin/python3
# ====================================================================
# FFT demo
#
# From: https://www.youtube.com/watch?v=O0Y8FChBaFU
# How to Compute FFT and Plot Frequency Spectrum in
# Python using Numpy and Matplotlib
# ====================================================================
import numpy as np
import matplotlib.pyplot as plt
# ---- signal frequency (Hz) (base frequency for this demo)
sig_freq = 100
# ---- sampling frequency (Hz)
sam_freq = 2000
# ---- time interval between two time samples
sam_tim_step = 1/sam_freq
# ---- number of samples
# ---- multiply by 10 to get a finer resolution in the
# ---- sampling of the signal
n = int(10*sam_freq/sig_freq)
# ---- frequency interval between two frequency samples
sam_freq_step = sam_freq/n
# ---- time steps (time series for sampling a signal)
# ---- t is a numpy array
t = np.linspace(0,(n-1)*sam_tim_step,n)
# ---- frequency steps (frequency series - one for each time step)
# ---- f is a numpy array
f = np.linspace(0,(n-1)*sam_freq_step,n)
# ----
# ---- in this demo the frequency spectrums are the same (1)
# ----
# ---- signal 1 - sine function - frequency 100
# ---- y1 is a numpy array
y1 = np.sin(2 * np.pi * sig_freq * t)
# ---- signal 2 - sine function - frequence 200
# ---- y2 is a numpy array
y2 = np.sin(2 * np.pi * sig_freq * t * 2)
# ---- signal 3 - sine function - frequence 50
# ---- y3 is a numpy array
y3 = np.sin(2 * np.pi * sig_freq * t * 0.5)
# ---- combine the signals (y is a numpy array)
y = y1 + y2 + y3
# ---- perform fft (x is a series of complex numbers)
x = np.fft.fft(y)
# ---- convert fft results to absolute values
# ---- normalize the data (divide by n)
x_abs = np.abs(x)/n
# ---- display environment
print()
print(f'signal frequency = {sig_freq}')
print(f'sample frequency = {sam_freq}')
print(f'sample time step = {sam_tim_step}')
print(f'sample freq step = {sam_freq_step}')
print(f't array len(t)={len(t)} t[0]={t[0]} t[-1]={t[-1]}')
print(f'f array len(f)={len(f)} f[0]={f[0]} f[-1]={f[-1] }')
print()
print(f'y1 array len(y1)={len(y1)} y1[0]={y1[0]} y1[-1]={y1[-1]:.6f}')
print(f'y2 array len(y2)={len(y2)} y2[0]={y2[0]} y2[-1]={y2[-1]:.6f}')
print(f'y3 array len(y3)={len(y3)} y3[0]={y3[0]} y3[-1]={y3[-1]:.6f}')
print(f'y array len(y) ={len(y)} y[0]={y[0]} y[-1] ={y[-1]:.6f}')
print()
print(f'x array len(x)={len(x)} x[0]={x[0]:.6f} x[-1]={x[-1]:.6f}')
print()
print(f'x_abs array len(x_abs)={len(x_abs)} x_abs[0]={x_abs[0]:.6f}' +\
f'x_abs[-1]={x_abs[-1]:.6f}')
print()
# ---- draw plots (four subplots) ------------------------------------
fig,[ax1,ax2,ax3,ax4] = plt.subplots(nrows=4, ncols=1,
figsize=(10,8),
constrained_layout=True)
fig.suptitle("FFT", fontsize=16)
# ---- plot the time domain
ax1.plot(f,y,'.-')
ax1.title.set_text('Time Domain')
ax1.set_ylabel('Amplitude')
ax1.set_xlabel('Time (sec)')
ax1.grid()
# ---- plot the frequency domain
ax2.plot(f,x_abs,'.-')
ax2.title.set_text('Frequency Domain')
ax2.set_ylabel('Amiplude')
ax2.set_xlabel('Frequencies')
ax2.grid()
# ---- plot half of the frequency domain
# ----
# ---- DC (direct current) (frequence 0) does not need
# ---- to be multiplied by 2. Otherwise, it will
# ---- overshadow the other values.
freq_half_plot = f[0:int(n/2+1)]
x_abs_half_plot = 2 * x_abs[0:int(n/2+1)]
x_abs_half_plot[0]= x_abs_half_plot[0]/2
ax3.plot(freq_half_plot,x_abs_half_plot,'.-')
ax3.title.set_text('Expanded Frequency Domain (1/2)')
ax3.set_ylabel('Amplitude')
ax3.set_xlabel('Frequency')
ax3.grid()
# ---- plot a quarter of the frequency domain
# ----
# ---- DC (direct current) (frequence 0) does not need
# ---- to be multiplied by 2. Otherwise, it will
# ---- overshadow the other values.
freq_quarter_plot = f[0:int(n/4+1)]
x_abs_quarter_plot = 4 * x_abs[0:int(n/4+1)]
x_abs_quarter_plot[0]= x_abs_quarter_plot[0]/4
ax4.plot(freq_quarter_plot,x_abs_quarter_plot,'.-')
ax4.title.set_text('Expanded Frequency Domain (1/4)')
ax4.set_ylabel('Amplitude')
ax4.set_xlabel('Frequency')
ax4.grid()
plt.tight_layout()
plt.show()