# animBall.py
# https://matplotlib.org/stable/gallery/animation/index.html
'''
  An animated trajectory of a cannon ball, which repeats.
  Clicking on the window causes the animation to pause/resume.
'''

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


def onClick(event):
  global animRunning
  if animRunning:
    anim.event_source.stop()
    animRunning = False
  else:
    anim.event_source.start()
    animRunning = True


def update(i):
  if i == 0:  # reset for next play
    ax.clear()
    ax.set(xlim=[0, 3], ylim=[-4, 10], 
           xlabel='Time', ylabel='Z') 
    ax.axhline(y=0, color='k', alpha=0.5)
  ax.scatter(ts[:i], zs[:i], c="b", s=10)


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

ts = linspace(0, 3, 40)
g = -9.81
v0 = 12
zs = [((g*t**2)/2 + v0*t) for t in ts]

fig, ax = plt.subplots()
fig.canvas.mpl_connect('button_press_event', onClick)

anim = animation.FuncAnimation(fig, update, 
               frames=len(ts)-1, interval=50) # ms
animRunning = True

plt.show()