#! /usr/bin/env python
########################################################################
#
# Import mathematical functions from the maths module, they are also
# available in numpy.  Note that the arguments in sin and cos are not in
# degrees but in radians
#
# Note: We do two ways of importing stuff here.  
# - the syntax "from module import function(s)" means that the code knows the
#   function(s) in question, without further reference to the module.
# - the syntax "import module" imports the module as a namespace, and when
#   you want to access its methods you must use the prefix "module."
#

from math import sin, cos, sqrt, pi, exp
import sys

import pylab

#
########################################################################
#
# Main calculation - some little helpers
#
# - calculate |x| = squareroot(x^2+y^2)
#

def quad_sum(x, y):
    return sqrt(x**2 + y**2)

#
# - a template for the change of the drag coefficient with air pressure/altitude
#   of the projectile
#

def drag_coefficient(b2_mass0,y,scheme):
    factor = 0.
    if scheme=='no':
        factor = 1.
    else:
        print "Air pressure scheme not know, return drag_coefficient = 0."
    return factor * b2_mass0

#
# - one single step in the Euler method:
#   Update the positions & velocities:
#      x_{i+1}   = x_i + v_{x,i} * (Delta t) 
#      y_{i+1}   = y_i + v_{y,i} * (Delta t)   
#      v_{x,i+1} = v_{x,i} - (B/m) * v_{x,i}v_i * (Delta t)
#      v_{y,i+1} = v_{y,i} - g * (Delta t) - (B/m) * v_{x,i}v_i * (Delta t)
#   Note: B/m is the parameter b2_mass.
#

def cannon_euler_step(xs, vxs, ys, vys, b2_mass0, g, dt):
    """Take lists of coordinate values and append a new value using the euler method."""
    #
    # Note: a[-1] denotes the last entry in list a
    # Also: We already use a template for the air-pressure dependence here
    # 
    b2_mass = drag_coefficient(b2_mass0,ys[-1],"no")

    v     = quad_sum(vxs[-1], vys[-1])
    newvx = vxs[-1] + dt*(-vxs[-1] * v * b2_mass)
    newvy = vys[-1] + dt*(-vys[-1] * v * b2_mass - g)
    newx  = xs[-1]  + dt*(vxs[-1])
    newy  = ys[-1]  + dt*(vys[-1])
    
    xs.append(newx)
    vxs.append(newvx)
    ys.append(newy)
    vys.append(newvy)

#
# - calculate the range of the projectile from the last y above and the first y
#   below 0:
#   x_range = (x_n y_{n+1}-y_n x_{n+1})/(y_{n+1}-y_n)
#

def find_impact(xs, ys):
    """Linear interpolation between the last points to find the impact location."""
    lastgrad = (ys[-1] - ys[-2])/(xs[-1] - xs[-2])
    x_impact = xs[-2] - (ys[-2]/lastgrad)
    ## reset final point to be impact point
    xs[-1] = x_impact
    ys[-1] = 0
    return x_impact

#
# - calculate the full trajectory in dependence on v_0, theta_0, b2/m and (Delta t)
#   this is done by repeating the euler step described above, until y<0.
#   For the homework assignment, you will have to add some arguments (for the way
#   the air density is calculated and for the stepper method).  
#

def calc_trajectory(launchvelocity, launchangle, b2_mass, dt):
    g = 9.81
    #
    # Initialise posn & velocity lists
    # 
    xs  = [0]
    ys  = [0]
    vxs = [launchvelocity * cos(pi*launchangle/180)]
    vys = [launchvelocity * sin(pi*launchangle/180)]

    # ys[-1] is the last entry in the list ys
    while ys[-1] >= 0:
        cannon_euler_step(xs, vxs, ys, vys, b2_mass, g, dt)

    # Interpolate (linearly) to find impact point
    x_impact = find_impact(xs,ys)

    print "Range of shot = %2.2f m" % x_impact 
    print "  for velocity = %2.2f m/s, angle = %2.2f deg and drag = %2.2e 1/m." \
        % (launchvelocity,launchangle,b2_mass)
    print "  Stepsize was %2.2f s." % dt
    return xs, ys

#
# A simple test program to plot a trajectory for different angles and stepsizes.
# Note the way the trajectories are caluclated and the legends are filled inside
# a for-loop.  
#

if __name__=="__main__":
    launchvelocity = 700.
    b2_mass        = 4.e-5
    launchangles   = [30.,45.,60.,75.]
    dts            = [0.1,1.,10.]

    pylab.subplot(2,1,1)
    pylab.title('Projectile trajectory with different angles, dt = 1s')
    pylab.ylabel('Height / m')
    pylab.setp(pylab.gca(), 'xticklabels', [])
    legend_plots, legend_names = [], []
    for launchangle in launchangles:
        xs, ys = calc_trajectory(launchvelocity, launchangle, b2_mass, dts[1])
        plot = pylab.plot(xs,ys, linewidth=1.0)
        legend_plots.append(plot)
        legend_names.append(str(launchangle)+' deg')
    pylab.legend(legend_plots, legend_names)
    pylab.grid(True, alpha=0.2)

    pylab.subplot(2,1,2)
    pylab.title('Projectile trajectory with different step sizes, angle = 50 deg')
    pylab.ylabel('Height / m')
    pylab.xlabel('Distance / m')
    legend_plots, legend_names = [], []
    for dt in dts:
        xs, ys = calc_trajectory(launchvelocity, launchangles[1], b2_mass, dt)
        plot = pylab.plot(xs,ys, linewidth=1.0)
        legend_plots.append(plot)
        legend_names.append(str(dt)+' s')
    pylab.legend(legend_plots, legend_names)
    pylab.grid(True, alpha=0.2)

    pylab.show()