import matplotlib.pyplot as plt
import numpy as np
import math
g=9.81
dt=1e-3
v0=40
angle=math.pi/4

#time axis
time=np.arange(0,10,dt)
def traj(angle,v0):
    vx0=math.cos(angle)*v0
    vy0=math.sin(angle)*v0
    #initialize x,y arrays
    x=np.zeros(len(time))
    y=np.zeros(len(time))
    # projectile starts at (0,0)
    x[0],y[0]=0,0
    #second elements of x and y
    x[1],y[1]=x[0]+vx0*dt,y[0]+vy0*dt
    i=1
    #continuous loop continuous until projectile hits ground
    while y[i] >=0:
        #fin x[i+1] and y[i+1]
        x[i+1]=(2*x[i]-x[i-1])
        y[i+1]=(2*y[i]-y[i-1])-g*dt**2
        #increments i for next iteration
        i=i+1
    #truncate x and y arrays
    x=x[0:i+1]
    y=y[0:i+1]
    #return x,y,flight time, range of projectile
    return x,y, (dt*i), x[i]
x, y, duration, distance = traj(angle,v0)
print( 'Distance:' ,distance)
print( 'Duration:' ,duration)

#part 2: range
n=40
x,y,duration,distance = traj(math.pi/4, 30)
velocities=np.linspace(0,40,num=n)
maxrange=np.zeros(n)
for i in range(n) :
    x,y,duration,maxrange[i] = traj(math.pi/4, velocities)
plt.plot(velocities,maxrange)
plt.xlabel("x")
plt.ylabel("y")
plt.xlim(0,400)
plt.ylim(0,360)
plt.show()