#  Demonstration code for Scientific Computing class,
#  http://www.math.nyu.edu/faculty/goodman/teaching/SciComp2022/
#  BlankSlate.py
#  uses Python 3 and Numpy
#  Python Classes Demo
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

Fint = 4    # number of time steps between plots
T    = 50.  # final time of the simulation
dt   = .01  # time step
m    = 4    # number of spirals

centersList = [ [ 0.,  0.],         # The centers of spirals
                [ 0.,  .2],         # as a Python list -- slow
                [ -.2, .1],
                [ .2,  .3]]
centers = np.array(centersList)     # centers info as numpy array

omegas     = np.array([ 1.,-.5,  .3,-.4])  # angular speeds of spirals
decayRates = np.array([ .1, .2, .05, 0.])  # decay rates (DUH!)

t = 0.
frameInfoList = []                 # data
while ( t < T):
   for step in range(Fint):
      xs = np.zeros(m)     # x coordinates of points
      ys = np.zeros(m)
      for k in range(m):
         center = centers[k,:]
         omega  = omegas[k]
         rate   = decayRates[k]
         L      = np.exp(-rate*t)
         xs[k]  = center[0] + L*np.cos(omega*t)
         ys[k]  = center[1] + L*np.sin(omega*t)
         t      = t + dt

   frameInfoList.append( [ xs, ys, t] )

fig, ax = plt.subplots()
ln, = ax.plot([], [], 'ro')
ax.set(aspect=1)
ax.set_title('spirals')
FrameSize = 1.5
time_text = ax.text(-.9*FrameSize, .9*FrameSize, 'hello')
time_text_framework = "time {t:9.2f}"
plt.grid()

def init():
    ax.set_xlim(-FrameSize, FrameSize)
    ax.set_ylim(-FrameSize, FrameSize)
    return ln,

def update(frame):
    xs = frame[0]
    ys = frame[1]
    t = frame[2]
    ln.set_data(xs,ys)
    title = "time is " + str(t)
    time_text.set_text(time_text_framework.format(t=t))
    
    return ln, time_text

ani = FuncAnimation(fig, update, frames = frameInfoList,
                    init_func=init, blit=True, interval = 30)
    
ani.save('movie.mp4')

plt.show()       # This crashes a lot on my laptop.  Dunno why