####################################################################
#
# The file and class for the first part of the fourth homework,
# dealing with motion in a central force, with absolute value given 
# by    F = kappa/r^eta.
#
# To solve the equations, we use kappa/m (kappa_by_m) and we use
# exponent for eta. 
#
# Note that the state vector for the test particle here is 
# x = np.array([x,y,vx,vy])
#
####################################################################

import numpy as np


class CentralForce:
    def __init__(self, kappa_by_m, exponent):
        """Initialize a central force system with m a = -kappa/r^exponent."""
        self.kappa_by_m = kappa_by_m
        self.exponent   = exponent
    
    def dx_dt(self,x,t):
        ### Fix this function ###
        return np.array([0,0,0,0])

def area(xvec1,xvec2):
    ### Fix this function ###
    return 0

############# a global helper ###############
        
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

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


if __name__ == "__main__":

    import pylab
    from Read_Type import read_float
    from DEq_Solver import DEq_Solver

    km = read_float('Kappa over m? ',0,default=1)
    exp = read_float('Exponent of (1/r) in the force? ',default=2)

    cf     = CentralForce(km,exp);
    solver = DEq_Solver(cf,'RK4')
    
    t0     = 0.
    t1     = 50.
    deltat = .01 
    x0     = np.array([0.,1.,-0.5,0.])
    history = solver.solve(x0,t0,t1,deltat)

    xs,ys = extractXY(history)
    ax = pylab.subplot('111')
    pylab.plot(xs,ys,'c',linewidth=1.0)
    pylab.plot([0],[0],'bx',[x0[0]],[x0[1]],'cx',linewidth=3)
    pylab.xlabel('x [m]')
    pylab.ylabel('y [m]')
    pylab.title('Central potential $a={%s}/{r^{%s}}$' % (km,exp))

    ax.set_aspect('equal')

    mid = len(history)/2
    print 'Start area   : %.6f' % area(history[0],history[1])
    print 'Midpoint area: %.6f' % area(history[mid],history[mid+1])
    print 'End area     : %.6f' % area(history[-2],history[-1])


    pylab.show()