####################################################################
#
# The file and class for the third homework, dealing with
# harmonic oscillation.
#
# The state vector here is x = np.array([phi,omega])
#
####################################################################

import numpy as np


class Pendulum:
    def __init__(self, L, q=0, F_D=0, Omega_D=0, nonlinear=False):
        """Initialize a pendulum object with 
        length L [m], 
        friction coefficient q [], 
        driving force amplitude F_D [N],
        frequency Omega_D [s^-1]. 
        and a switch to enable the nonlinear pendulum"""
        ##############################################################
        ### Set the parameters here
        ##############################################################
 

    def dx_dt(self,x,t):
        ##############################################################
        ### Implement the differential equation for the pendulum with
        ### all options: friction, driving force and switch between
        ### linear and nonlinear
        ##############################################################
        return np.array([0.1, 0])

    def dx_dt_EC(self,x,t,deltat):
        ##############################################################
        ### Re-implement dx_dt from above using the EulerCromer
        ### method. The DEq_Solver provides 'deltat' for you.
        ##############################################################
        return np.array([0.1, 0])

    def T(self,x,t):
        ##############################################################
        ### Return the kinetic energy of state x,t
        ##############################################################
        return 0

    def U(self,x,t):
        ##############################################################
        ### Return the potential energy of state x,t
        ### Use the 'nonlin' switch to choose between the exact and 
        ### linearized form.
        ##############################################################
        return 0
        

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

def extractTPhiOmega(history):
    """Extracts a list of phi and omega values from the solver history."""
    ts = []
    phis = []
    omegas = []
    for t, (x,y) in history:
        ts.append(t)
        phis.append(180.*x/np.pi)
        omegas.append(y)
    return ts,phis,omegas

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


if __name__ == "__main__":

    import pylab

    from Read_Type import read_float
    from DEq_Solver import DEq_Solver

    phi = read_float('Initial angle (deg)? ', 0)
    phi *= np.pi / 180.
    x0 = np.array([phi,0])

    t0 = 0.
 
    L = read_float('Length (m)? ', 0)
    q = read_float('Drag coeffcient q (s^-1)? ',0)

    FD = read_float('Driving force amplitude (s^-2)? ', 0)
    OD = read_float('Driving force frequency (s^-1)? ', 0)

    t1 = read_float('End time t1 (s)?', 0)

    deltat = 0.02

    method = 'Euler' 

    ####### Switch to 'EulerCromer' to test that part of the code
    #method = 'EulerCromer'
    

    label = { True : 'nonlinear', False : 'linear' }
    style = { True : 'b-.', False : 'r-' }

    pylab.subplot(211)

    pylab.axhline(0,color='lightgrey')

    for nonlin in [ True, False ]:
        pend = Pendulum(L,q,FD,OD,nonlin)
        solver = DEq_Solver(pend, method)
        result = solver.solve(x0,t0,t1,0.02,circular_index=0)

        ts,xs,_ = extractTPhiOmega(result)    

        pylab.plot(ts,xs,style[nonlin],label=label[nonlin])


    pylab.legend()

    pylab.title('Pendulum. $L$=%s m; $q$=%s/s; $F_D$=%s$/\mathsf{s}^2$, $\Omega_D$=%s/s'
                %(L,q,FD,OD))
    pylab.ylabel('Angle [deg]')



    Ts = [ pend.T(x,t) for t,x in result ]
    Us = [ pend.U(x,t) for t,x in result ]
    Es = [ a + b for a,b in zip(Ts,Us) ]

    pylab.subplot(212)
    pylab.plot(ts,Ts,'b',label='$E_{kin}$')
    pylab.plot(ts,Us,'g',label='$E_{pot}$')
    pylab.plot(ts,Es,'r',label='$E_{tot}$')

    
    pylab.title('Energy balance, %s case. Method: %s' 
                % (label[nonlin], method) )
    pylab.xlabel('Time [s]')
    pylab.ylabel('Energy [J/kg]')
    pylab.legend()


    pylab.show()