# lorenzAnim.py
"""
  An animated version of the Lorenz model which can be 
  paused and resumed by clicking in the window. The 
  animation consists of a line representing the curve
  and also a single red dot to highlight the current position.
"""

from matplotlib import pyplot as plt
from matplotlib import animation
import math
from frange import *

DT = 0.01
N = 5000

def lorenz(xyz, s=10, p=28, b=8/3):
  x, y, z = xyz
  x_dot = s*(y - x)
  y_dot = p*x - y - x*z
  z_dot = x*y - b*z
  return (x_dot, y_dot, z_dot)


def update(i):
  # update the curve and the red dot
  global currFrame
  line.set_data(xs[:i], ys[:i])
  line.set_3d_properties(zs[:i])
  scat._offsets3d = (xs[i:i+1], ys[i:i+1], zs[i:i+1])
  currFrame = i
  if currFrame == N-1:
    print("Finished")


def onClick(event):
  # pause/resume the animation
  global animRunning
  if currFrame < N-1:
    if animRunning:
      anim.event_source.stop()
      animRunning = False
      print(f"Stopped at frame {currFrame} at ({xs[currFrame]:.3f}, {ys[currFrame]:.3f}, {zs[currFrame]:.3f})")
    else:
      anim.event_source.start()
      animRunning = True
      print("Resumed...")


# ------------------ main ----------------

xyz = [(0,0,0) for x in range(N+1)]
       # Need one more for the initial values
xyz[0] = (0, 1.0, 1.05)

for i in range(N):
    xyz[i+1] = tuple(map(lambda x,y: x+y*DT, 
                                xyz[i], lorenz(xyz[i])))
                # xyz[i] + lorenz(xyz[i]) * DT

xs, ys, zs = zip(*xyz)

fig = plt.figure(figsize=(8, 8))
fig.canvas.mpl_connect('button_press_event', onClick)
ax = fig.add_subplot(projection='3d')

line, = ax.plot(xs[0:1], ys[0:1], zs[0:1])
scat = ax.scatter(xs[0:1], ys[0:1], zs[0:1], color="r", s = 50)

ax.set_xlim3d([-20.0, 20.0])
ax.set_xlabel('X')
ax.set_ylim3d([-20.0, 20.0])
ax.set_ylabel('Y')
ax.set_zlim3d([0.0, 50.0])
ax.set_zlabel('Z')
ax.set_title("Animated Lorenz Attractor")

anim = animation.FuncAnimation(fig, update, repeat=False,
                      frames=N, interval=10)   # ms
animRunning = True
currFrame = 0

plt.show()