import numpy as np
import scipy.fftpack as fft
import matplotlib.pyplot as plt

## genero la senal espacial
N = 2**12
x = np.linspace(0, 20*np.pi, N)
dx = x[1] - x[0]

y = np.sin(x) + 0.5*np.sin(2*x) + 0.8*np.sin(4.5*x) + 0.1

## calculo la transformada de fourier
Fy = fft.fft(y)
# genero vector con las frecuencias
# para el caso temporal fftfreq es la frecuencia f;
# omega (o k para el caso espacial) va multiplicada por 2*pi
k = 2 * np.pi * fft.fftfreq(N, dx)

# reordeno los vectores para que esten en orden de frecuencias creciente
# (no es necesario para graficar abajo; lo agrego por completitud)
Fy_ord = fft.fftshift(Fy)
k_ord = fft.fftshift(k)

## reconstruyo la senal original invirtiendo la transformada de fourier
yc = fft.ifft(Fy)
# los valores de yc son complejos, pero sabemos que la senal es real
# (impriman los valores de yc en pantalla y van a ver que las partes
#  imaginarias son practicamente 0; i.e. terminos del order de 10^-17)
yr = np.real(yc)

## grafico los resultados
fig = plt.figure(figsize=(12, 6))

# espacial
ax1 = fig.add_subplot(211)
ax1.plot(x, y, '-', color='royalblue', label='original')
ax1.plot(x, yr, '--', color='orangered', label='reconstruida')
ax1.grid()
ax1.set_xlabel('$x$')
ax1.set_ylabel('$y(x)$')
ax1.legend()
ax1.set_title('Prueba transformada Fourier')

# fourier
ax2 = fig.add_subplot(212)
ax2.plot(k_ord, np.abs(Fy_ord)**2, '-o', color='forestgreen', markersize=2, label='espectro')
ax2.grid()
ax2.set_xlabel('$k$')
ax2.set_ylabel('$|\mathcal{F}(y)|^2$')
ax2.legend()

fig.tight_layout()
plt.show()