"""
Exercise 5.10
This is the same as 5.9, except we will now read
multiple v0 values from the command line, and plot
the curve for all.
"""

import numpy as np
import matplotlib.pyplot as plt
import sys

"""
We use a list comprehension to convert the
list of command line arguments to a list of
floats:
"""
v0_list = [float(v) for v in sys.argv[1:]]

g = 9.81

#loop over v0 values, compute new t,y arrays and plot:
for v0 in v0_list:
    t_stop = 2*v0/g
    t = np.linspace(0,t_stop,101)
    y = v0*t -0.5*g*t**2
    plt.plot(t,y,label='v0 = %g' %(v0))

#decorate and show plot
plt.xlabel('time (s)')
plt.ylabel('height (m)')
plt.title('Height of ball')
plt.legend() #legend based on the labels defined above

plt.show()

"""
Terminal> python plot_ball2.py 5 10 15
(output is a plot)
"""