#  Python 3
#  Author: Jonathan Goodman, goodman@cims.nyu.edu
#  Fall 2024
#  for the class Scientific Computing
#  https://math.nyu.edu/~goodman/teaching/ScientificComputing2024/ScientificComputing.html

#  Movie and color demo for Assignment 5

# This file: colorPlay.py,
# Writes a movie file: see the line: movieFileName = ...

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.animation import FFMpegWriter

phi_min = 0.          #  range of phi angle
phi_max = 2*np.pi
om      = 3.          #  angular speed of the color wave
lam     = .3          #  decay rate for color wave
A       = 50.         #  max amplitude of the color wave

n  = 220                          # number of points in a plot curve
nf = 100                          # number of movie frames
dp = ( phi_max - phi_min)/nf      # phi difference between frames

movieFileName = "colorWaveMovie.mp4"

#  Put points on a fixed curve with colors representing the intensities
#  The frames have the same point coordinates but different colors

#  Calculate the point coordinates, a full period of the sine curve.

th = np.linspace(0., 2*np.pi, n)   # th = theta, evenly spaced in [0,2*pi]
y  = np.sin(th)

#   create the values to be visualized with color
#   v[fr, k] will be the intensity at point k in frame number fr

v = np.zeros([nf,n])
for fr in range(nf):           # fr = frame number
   phi = phi_min + dp*fr       # phi angle for this frame
   for k in range(n):
   
           # color function = rotating and decaying wave
    
      v[fr,k]= A*np.cos( th[k] - om*phi) * np.exp(-lam*phi)

v_max = np.max(v)    #   the range of values is used to set the color scale
v_min = np.min(v)

fig,ax = plt.subplots()

cmap = mpl.cm.hot    #   a choice of colormap.  The "colormap" documentation
                     #   See the "colormap" documentation in matplotlib
                     
#  Layout the plot, with the picture on the left and colorbar on the right

ax.scatter( th, y, c=v[0,:], cmap = cmap)
fig.colorbar(mpl.cm.ScalarMappable(norm=mpl.colors.Normalize(v_min, v_max),
            cmap=cmap), ax=ax, orientation='vertical',
            label='brightness scale')

ax.grid()
plt.rcParams['text.usetex'] = True
ax.set_title(r"phase angle $\phi = 0$")
ax.set_ylabel(r"$y$")
ax.set_xlabel(r'$\theta$')

mpl.use("Agg")              # I don't know what this does

metadata = dict(title='color wave movie', artist='Matplotlib',
                comment='Do not know what I am doing!')
writer = FFMpegWriter(fps=15, metadata=metadata)

with writer.saving(fig, movieFileName, 100):   # 100 is the dpi resolution
    for fr in range(nf):
        
        phi = phi_min + dp*fr
        ax.set_title(r"phase angle $\phi = {phi:5.2f}$".format(phi=phi))

        ax.scatter( th, y, c=v[fr,:], cmap = cmap,
                    vmin = v_min, vmax = v_max)
        
        writer.grab_frame()
print("All done, movie saved")