####################################################################
#
# The file and class for the second part of the fourth homework,
# dealing with planetary movement in a two-planet system, i.e.
# a three-body problem.  We call these three bodies Sun, Earth, and
# Jupiter.  To initialise the problem, the three masses (in kg),
# the initial x positions (in m, note that all objects have y=0 at 
# start time), and the x- and y-components of the Earth's and the 
# Jupiter's velocity (in m/s) are necessary.  The initial velocity 
# of the sun is determined by momentum conservation, to allow for 
# nicer plots. 
#
# The state vector here is 
# x = np.array([xS,yS,vxS,vyS,xE,yY,vxE,vyE,xJ,yJ,vxJ,vyJ])
#
####################################################################

import numpy as np


class SolarSystem:
    def __init__(self, mass_sun, mass_earth, mass_jupiter):
        """Initialize a solar system with mass of the sun [kg],
        earth [kg] and jupiter[kg]."""
        self.M     = float(mass_sun)
        self.mE    = float(mass_earth)
        self.mJ    = float(mass_jupiter)
        self.G     = 6.67428e-11
        print 'Initialised solar system with MS = ',self.M,' and G = ',self.G,'.'
        print '   Planetary masses: earth = ',self.mE,' and jupiter = ',self.mJ,'.'

    def dx_dt(self,x,t):
        ### Fix this function ###
        return np.array([0,0,0,0,0,0,0,0,0,0,0,0])


############# a global helper ###############
        
def extract3XY(history):
    """Extracts a list of x and y values from the solver history."""
    xSs = []
    ySs = []
    xEs = []
    yEs = []
    xJs = []
    yJs = []
    for _, (xS,yS,vxS,vyS,xE,yE,vxE,vyE,xJ,yJ,vxJ,vyJ) in history:
        xSs.append(xS)
        ySs.append(yS)
        xEs.append(xE)
        yEs.append(yE)
        xJs.append(xJ)
        yJs.append(yJ)
    return xSs,ySs,xEs,yEs,xJs,yJs

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


if __name__ == "__main__":

    import pylab
    from Read_Type import read_float
    from DEq_Solver import DEq_Solver

    M       = read_float('Mass of the sun? ', 0.,default=2.0e+30)
    mE      = read_float('Mass of the earth? ', 0.,default=6.0e+24)
    mJ      = read_float('Mass of jupiter? ', 0.,default=1.9e+27)
    planets = SolarSystem(M,mE,mJ)

    x0E = read_float('Initial x-distance of earth? ',default=1.5e+11)
    vxE = read_float('Initial x-velocity of earth? ',default=0)
    vyE = read_float('Initial y-velocity of earth? ',default=3.0e+4)
    x0J = read_float('Initial x-distance of jupiter? ', default=-7.78e+11)
    vxJ = read_float('Initial x-velocity of jupiter? ', default=0)
    vyJ = read_float('Initial y-velocity of jupiter? ', default=-1.3e+4)
    vxS = -(vxE*mE+vxJ*mJ)/M
    vyS = -(vyE*mE+vyJ*mJ)/M
    x      = np.array([0.,0.,vxS,vyS,x0E,0.,vxE,vyE,x0J,0.,vxJ,vyJ])
    t0     = 0
    t1     = read_float('End time t1? ',0.,default=4e7)
    deltat = read_float('Timestep? ',1.e4,default=1e4)

    print 'Calculate trajectories of the planets starting with:'
    print '   earth   = ',x0E,' v = (',vxE,', ',vyE,')'
    print '   jupiter = ',x0J,' v = (',vxJ,', ',vyJ,')'

    solver = DEq_Solver(planets, 'RK4')
    result = solver.solve(x,t0,t1,deltat)

    xSs,ySs,xEs,yEs,xJs,yJs  = extract3XY(result)    

    pylab.ion()

    ax = pylab.subplot('111')

    pylab.plot(xJs[:1],yJs[:1],'x',color='orange',linewidth=1.0)
    pylab.plot(xEs[:1],yEs[:1],'xb',linewidth=1.0)
    pylab.plot(xSs[:1],ySs[:1],'xy',linewidth=2.0)
   # Set up first subplot
    trajectoryJ, = pylab.plot(xJs[:1],yJs[:1],color='orange',linewidth=1.0)
    trajectoryE, = pylab.plot(xEs[:1],yEs[:1],'b',linewidth=1.0)
    trajectoryS, = pylab.plot(xSs[:1],ySs[:1],'y',linewidth=2.0)

    posJ, = pylab.plot(xJs[:1],yJs[:1],'o',color='orange',linewidth=10.0)
    posE, = pylab.plot(xEs[:1],yEs[:1],'bo',linewidth=10.0)
    posS, = pylab.plot(xSs[:1],ySs[:1],'yo',linewidth=10.0)



    # Annotate the graph
    pylab.title('Planetary movement')
    pylab.xlabel('x [m]')
    pylab.ylabel('y [m]')
    pylab.legend((trajectoryS,trajectoryJ,trajectoryE),('Sun %s'%M,'Jupiter %s'%mJ,'Earth %s'%mE))
    pylab.grid(True)

    data = xJs + xEs + xSs + yJs + yEs + ySs
    mind = min(data)
    maxd = max(data)

    limit = 1.05 * max(abs(mind),abs(maxd))

    pylab.xlim(-limit,limit)
    pylab.ylim(-limit,limit)
    ax.set_aspect('equal')

    step = len(result)/200

    for i in range(0,len(result)-step,step):
        pylab.title('Planetary movement %s %%' % (int(100*i)/int(len(result))))
        trajectoryJ.set_data(xJs[i:i+step], yJs[i:i+step])
        trajectoryE.set_data(xEs[i:i+step], yEs[i:i+step])
        trajectoryS.set_data(xSs[i:i+step], ySs[i:i+step])
        posJ.set_data(xJs[i+step],yJs[i+step])
        posE.set_data(xEs[i+step],yEs[i+step])
        posS.set_data(xSs[i+step],ySs[i+step])
        pylab.draw()

    pylab.title('Planetary movement 100 %')
    trajectoryJ.set_data(xJs, yJs)
    trajectoryE.set_data(xEs, yEs)
    trajectoryS.set_data(xSs, ySs)
    pylab.draw()

    pylab.ioff()

    pylab.show()