# -*- coding: utf-8 -*-
"""
Created on Sun Feb  7 20:04:31 2021
2D ejecta problem, including air resistance proportional to v**2
Note careful math for the derivative function, since the non-linearity
creates a sign problem when taking the derivative. Careful use of alternative
arctan2() function would make it shorter, but maybe less obvious
@author: neil
"""
import numpy as np
import matplotlib.pyplot as plt

# ************HEADER****************************
cons   = .02
grav   = 9.8
tarray = np.linspace(0,7,50)   # a list of times in the solution  
v0     = [20,30]               # initial conditions for the vx,vz velocities m/s
# vx must be positive, vz can be up or down

# **This derivative is tricky since it is non-linear and v*v removes the sign of v***
def myderiv(v,t):   # returns the vector components of the deriv of v at t
                    # uses gobals cons and grav, v[0] - Vx, v[1] - Vy
    if len(v) != 2: # must have input velocities for both x and z
        print('Input error, v[] must be length 2')
        return [0,0]
    # the signs are tricky, need to make sure we are in the right hand quadrant
    vel   = np.sqrt(v[0]*v[0]+v[1]*v[1])
    drag  = cons*vel*vel        # need to make positive 
    if v[0] == 0:               # avoid a divide by 0 error, theta is velocity vector angle
        theta = np.pi/2         # theta is needed to partition drag between x,z
    else:                       # ensure this never goes into negative since
        theta = np.arctan(v[1]/v[0])  # probably could avoid all this sign trouble
    sign = -1                         # using np.arctan2() instead ??
    if v[1] < 0:
        sign = 1
    theta = np.abs(theta)   # make sure theta is in the right hand quadrant
                            # pi/2 to - pi/2

    dvx   = -drag*np.cos(theta) # assume vx always positive, so drag is negative
    dvz   = -grav + sign*drag*np.sin(theta)
    return dvx,dvz
        
# ***************** Make our own ODE method, note it doesn't have any checks, nor is it very accurate
def myodeint(f,vstart,t): # make our own ode method
    """ primative ODE (forward difference), takes: a derivative function, start point, and
    evenly spaced vector of solution points"""
    v = np.zeros([len(t),2])   # make an array to output our results of velocities
    delt = t[1] - t[0]         # assume evenly spaced solution points, so delt is the spacing
    v[0,0] = vstart[0]         # first output is the starting point
    v[0,1] = vstart[1]
    
    for n in np.arange(1,len(t)):       #fill in all the other points after the first
        dvx,dvy = f(v[n-1,:],1)         # get the deriv of both vx,vz,  1 is a dummy argument
        v[n,0] = v[n-1,0]+dvx*delt      #this ia a standard Euler forward difference representation
        v[n,1] = v[n-1,1]+dvy*delt
    return v                            # return our calculated solution vector at the 't' points

Vz  = myodeint(myderiv,v0,tarray)     # myderiv is our function, Vz is an array of velocities vs time 

# ************************that is the entire program*******************************************
# for plotting, convert t to x using dz/dt = v, or dz = v*dt
dt = tarray[1]-tarray[0]
z = np.empty_like(tarray)  # create an array the same length as tarray
x = np.empty_like(tarray)  # create an array the same length as tarray
z[0] = 0                   # set the first value to 0
x[0] = 0
for n in np.arange(1,len(tarray)): # step thru Vz, summing the distance traveled at each time step
    z[n] = z[n-1] + dt*Vz[n,1]
    x[n] = x[n-1] + dt*Vz[n,0]

# this shows making a multi-panel plot to show various aspects of the solution
# make our basic plotting canvas, with a title in the header bar
fig = plt.figure()
fig.canvas.set_window_title ("Plots for falling bodies, using a coefficent of {:6.3f}".format(cons))

# Make a 2x2 array of plots, and start filling the first one
ax1 = fig.add_subplot(1,2,1)
ax1.grid(True)
ax1.plot(tarray,Vz[:,1], 'r-+')  
ax1.set_title("Velocity vs Time for a falling body")
ax1.set_xlabel("time (secs)")  
ax1.set_ylabel("velocity (m/s)")    

# now we have x as well as t, we can plot the other 2 views
ax2 = fig.add_subplot(1,2,2)
ax2.grid(True)
ax2.plot(x,z, 'b-+')  
ax2.set_title("Trajectory profile for a falling body")
ax2.set_xlabel("X (m)")  
ax2.set_ylabel("Z (m)")