"""
Ex 5.11 from "A primer on..."
Same as 5.10, but this time compute the
max t and y values and use these to set
the length of the axes in the plot. 
"""

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

v0_list = sys.argv[1:]

g = 9.81

max_t = 0
max_y = 0

for v0 in v0_list:
    v0 = float(v0)
    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 = f'v0 = {v0}')

    if max(t) > max_t:
        max_t = max(t)
    if max(y) > max_y:
        max_y = max(y)


plt.axis([0, max_t, 0, max_y * 1.05])
plt.legend()
plt.xlabel('time (s)')
plt.ylabel('height (m)')
plt.show()