# P372-PS14-FFT
# Matthew Oros
"""
This script introduces the Real FFT and Inverse Real FFT capabilities.
"""

import numpy as np
import pylab as pl
import numpy.fft as fft

# Some constants, I left these all the same
numPoints = 1000  # number of points in x-waveform, how much resolution
numFreqs = numPoints / 2 + 1  # (numPoints/2 -1) complex & 2 real k-amplitudes
# This is still same numPoints of bits of information!
# The wavenumber k's are also called "spatial frequencies"
# We are going to look at the region between 0 and 2 in x-space,
# but suppose we don't want anything to repeat within twice this, so
windowLength = 4.0  # effective window length
dx = windowLength / numPoints  # Step size in x-space
x = np.arange(0.0, windowLength, dx)  # All the x-positions

# Now prepare the spectrum and k wavenumbers
dk = 2 * np.pi / windowLength  # Step size in k-space
ks = np.arange(numFreqs) * dk  # All the k's, highest k found at ks[-1]

# HERE IS WHERE THE REAL WORK GETS DONE, the above is all just setting up
# Choose the <k>
k0 = 0.1 * ks[-1]  # I chose 0.1 to have lower frequency inside the wave packet
alpha = 20  # 0.1  # I chose 0.1 to make the wave packet wider
# Now define the k-spectrum amplitudes
# phi_k = np.exp(-0.5 * alpha * (ks - k0) ** 2)  # the spectral amplitudes (complex)
# Upside-down parabola!
phi_k = np.piecewise(ks,
                     [(ks - k0 < -alpha), (-alpha <= ks - k0) & (ks - k0 <= alpha), (ks - k0 > alpha)],
                     [0, lambda k: (1 / alpha ** 2) * (k - k0) ** 2, 0])

# And generate the wavefunction from phi(k) using Inverse Real FFT
psi_x = fft.irfft(phi_k)  # inverse real FFT gets the real wavefunction

# I will now normalize the wave function,
# which was not done in the original script
norm_factor = 1 / np.sqrt(np.sum(psi_x * psi_x))
psi_x *= norm_factor

pl.figure(1)
pl.plot(ks, phi_k)  # phi_k could be complex, which would throw a warning here
# if so, could use: pl.plot(ks, np.abs(phi_k) )
pl.title(r"Spectrum $\phi(k)$")
pl.xlabel("wavenumber k")
pl.ylabel("spectral amplitude")

pl.figure(2)
pl.plot(x, psi_x)
pl.title(r"Wavefunction $\psi(x)$")
pl.xlabel("position x")
pl.ylabel(r"wavefunction $\psi(x)$")
pl.annotate("This piece also appears to left of x=0!", (3.8, -0.010),
            xytext=(1.0, -0.03), arrowprops=dict(facecolor='black', shrink=0.05))

# Here's a fancier figure, without recalculating anything but noticing
# that the wavepacket right now is centered around x=0.  Some of that is
# showing up on the right side of the window, because the whole signal is
# repeating with a period of 1 windowLength, so we will copy that and stitch
# it on the left side in a different color.  (If you look closely, you will
# see that the last point of the "orange" needs to be connected to the first
# point of the "blue")
# In the NEXT script P372-PS14-FFT-shift.py, the wavepacket is shifted in x-space
pl.figure(3)
cutPoint = (3 * numPoints) // 4  # 3/4 of the way through the window
pl.plot(x[0:cutPoint], psi_x[0:cutPoint])  # as before, but only to cutPoint
pl.plot(x[cutPoint:] - windowLength, psi_x[cutPoint:])  # put last 1/4 on left side
pl.plot([-dx, 0], [psi_x[-1], psi_x[0]], 'k')  # Final stitch
pl.title(r"Wavefunction $\psi(x)$")
pl.xlabel("position x")
pl.ylabel(r"wavefunction $\psi(x)$")

pl.show()  # show the graph windows and loop here until they are closed

print("Done.")