# P372-PS14-FFT-animate.py
"""
A wavepacket is made to move.  Based on P372-PS14-FFT.py
and uses matplotlib.animation library
"""

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

# Some basic numbers, first.  These are all design decisions you can change!
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]  # 25% of highest k value <-- PICK THIS!
alpha = 0.1  # smaller alpha means wider phi_k, narrower psi_x <-- AND THIS!
# Now define the k-spectrum amplitudes
phi_k = np.exp(-0.5 * alpha * (ks - k0) ** 2)  # the spectral amplitudes (complex)
# And generate the wavefunction from phi(k) using Inverse Real FFT
psi_x = fft.irfft(phi_k)  # inverse real FFT gets the real wavefunction

# Now we want the wavepacket to move!  Let's define a velocity:
v = 1.0

fig, ax = pl.subplots()  # Create the graphics canvas
pl.title(r"Wavefunction $\Psi(x,t)$ animated")
pl.xlabel("position x")
pl.ylabel(r"wavefunction $\Psi(x,t)$")

T = 1.0  # total time duration
Nsteps = 200  # number of steps of animation
times = np.linspace(0, T, Nsteps + 1)  # the times at which to calculate the shifts
# Draw the first plot and get the lines object
lines = ax.plot(x[:numPoints // 2], psi_x[:numPoints // 2], lw=1)[0]


def update(frame):
    # This calculates the new values for the plot, called by the animation
    time_now = times[frame]
    # Velocity changes with time!
    v = np.cos(time_now * 10)
    # Calculate the phase shifts
    phase_shift = np.exp(-1j * ks * v * time_now)
    # And apply them before doing the F.T.
    psi_x_shift = fft.irfft(phi_k * phase_shift)
    lines.set_ydata(psi_x_shift[:numPoints // 2])
    return lines


# Connect the update function to the animation
ani = animation.FuncAnimation(fig=fig, func=update, frames=Nsteps, interval=30)

# Show the animation and run it until the window is closed
pl.show()

print("Done.")