####################################################################
#
# The file and class for the second homework, dealing with
# particle motion in one dimension.
#
# The second order differential equation can be rewritten as a first
# order equation using the variable x = numpy.array([x,y,vx,vy]). This
# allows the stepper to use dx_dt in the same way as last time.
#
####################################################################

import numpy as np


class Cannonball:
    def __init__(self,B_m,airmodel='const'):
        """Initialize a cannonball object with drag coefficient B/m
        and an optional air density model"""
        self.B_m = float(B_m)
        self.airmodel = airmodel

    def dx_dt(self,x,t):
        ####### FIX THIS FUNCTION (Q.1a) #######
        ### USE correctionFactor HERE (Q.1b) ###
        return np.array([1,1,0,0])

    def correctionFactor(self,height):
        """The air density correction rho/rho_0"""
        ### FIX THIS FUNCTION (Q.1b) ###
        if   self.airmodel == 'const':      return 1.0
        elif self.airmodel == 'isothermal': return 1.0
        elif self.airmodel == 'adiabatic':  return 1.0
        else:
            raise Exception('Unknown air model',self.airmodel)


############ global helper functions ############


def extractXY(history):
    """Extracts a list of x and y values from the solver history."""
    xs = []
    ys = []
    for _, (x,y,_,_) in history:
        xs.append(x)
        ys.append(y)
    return xs,ys



def range(history):
    """Calculates the range of the shot by interpolating between the
    last two points in the solver history."""
    from Interpolator import Interpolator
    xs, ys = extractXY(history)
    ### RETURN INTERPOLATED RANGE VALUE HERE (Q.3c) ###
    return 0




############# testing code ###############


if __name__ == "__main__":

    import pylab

    shapes  = {'const':'--','isothermal':'-.','adiabatic':'-'}
    colours = {'Euler':'y','RK2':'k','RK4':'r'}

    from Read_Type import read_float
    from DEq_Solver import DEq_Solver

    v0 = read_float('Initial velocity (m/s)? ', 0)
    phi = read_float('Initial angle (deg)? ', 0)
    phi *= np.pi / 180.

    x0 = np.array([0, 0, v0*np.cos(phi), v0*np.sin(phi)])

    t0 = 0.
    t1 = 100.

    B_m = read_float('Drag coeffcient B/m (m^-1, realistic value is 4e-5)? ',0)
    deltat = read_float('Size of timestep (s)? ', 0)



    ax1 = pylab.subplot(111)
    pylab.title('Cannon shot. v0=%s m/s, B/m=%s m^-1' % (v0,B_m))

    for airmodel in ['const','isothermal','adiabatic']:
        print 'Air model is:',airmodel
        cannon = Cannonball(B_m,airmodel)

        for method in ['Euler','RK2','RK4']:
            solver = DEq_Solver(cannon, method)
            result = solver.solve(x0,t0,t1,deltat,positive_index=1)

            print 'Method %s:\tRange=%.6f m' % (method,range(result))

            xs,ys = extractXY(result)    
            ax1.plot(xs,ys,shapes[airmodel]+colours[method],
                       label='%s %s'%(method,airmodel))

    # Annotate the graph
    ax1.set_ylabel('Height (m)')
    ax1.legend(loc=2)
    ax1.set_xlabel('Distance (m)')
    ax1.set_ylim(0)

    pylab.show()