import matplotlib.pyplot as plt
import numpy as np
from numpy.fft import fft, fftshift

plt.clf()

fs = 10  # sampling frequency

dt = 1 / fs

t = np.arange(0, 5, dt)  # start, end, step
y = np.cos(2 * np.pi * 1 * t)

t1 = np.arange(0, 5, dt / 1000)  # start, end, step
y1 = np.cos(2 * np.pi * 9 * t1)

N = len(y)

fig, axs = plt.subplots(2, num=1)

axs[0].plot(t, y, ':')
axs[0].stem(t, y, '.')
axs[0].set_xlabel('t [s]')
axs[0].set_ylabel('y')
axs[0].grid()

axs[0].plot(t1, y1, 'r')

Y = fftshift(fft(y) / N)
df = fs / len(Y)
f = np.arange(-fs / 2, fs / 2, df)

axs[1].stem(f, np.abs(Y))
axs[1].set_xlabel('f [Hz]')
axs[1].set_ylabel('Y abs')
axs[1].grid()

# axs[1].stem(f, Y.real)
# axs[1].set_xlabel('f [Hz]')
# axs[1].set_ylabel('Y real')
# axs[1].grid()

# axs[2].stem(f, Y.imag)
# axs[2].set_xlabel('f [Hz]')
# axs[2].set_ylabel('Y imag')
# axs[2].grid()

plt.show()