from numpy import array, cos, sin, shape, min, double, zeros, pi
import matplotlib.pyplot as plt
from ode45 import ode45
import scipy.optimize


def PendulumODE(y, l, g):
    """PendulumODE return the right-hand side of the math. pendulum"""
    dydt = array([y[1], -g * sin(y[0]) / l])
    return dydt

def rhs(y, l=0.6, g=9.81):
    """Right-hand side of the math. pendulum for Verlet methods"""
    return -g * sin(y) / l

def IntegratePendulum(phi0,phi1, tEnd=1.8, l=0.6, g=9.81, flag=False):
    """IntegratePendulum solve the mathematical pendulum with ode45
    flag == false   -->  y = complete solution phi,phi',t
    flag == true    -->  y = phi(tEnd)
    (order of the 2 outputs reversed wrt the usual to use fzero)
    """
    if shape(phi0) == (1,):
        phi0 = phi0[0]
        phi1 = phi1[0]
    y0 = array([phi0, phi1])
    tspan = (0, tEnd)
    t, y = ode45(lambda t, y: PendulumODE(y, l, g), tspan, y0)
    if flag:
        return y[-1][0]
    else:
        return (y, t)


def IntegratePendulumStVerlet2step(phi0,phi1, N, tEnd=1.8, l=0.6, g=9.81, flag=False):
    """IntegratePendulum solve the mathematical pendulum with two step Stoermer Verlet
    flag == false   -->  y = complete solution phi,phi',t
    flag == true    -->  y = phi(tEnd)
    """
    if shape(phi0) == (1,):
        phi0 = phi0[0]
        phi1 = phi1[0]
    y0 = array([phi0, phi1])

    # stepsize
    h = double(tEnd) / N
    # memory allocation
    y = zeros((N + 1, 2))
    y[0, :] = y0
    t = zeros(N + 1)

    #########################################
    #                                       #
    # Implement 2 step Stoermer Verlet here #
    #                                       #
    #########################################

    if flag:
        return y[-1][0]
    else:
        return (y, t)


def IntegratePendulumStVerlet1step(phi0,phi1, N, tEnd=1.8, l=0.6, g=9.81, flag=False):
    """IntegratePendulum solve the mathematical pendulum with two step Stoermer Verlet
    flag == false   -->  y = complete solution phi,phi',t
    flag == true    -->  y = phi(tEnd)
    """
    if shape(phi0) == (1,):
        phi0 = phi0[0]
        phi1 = phi1[0]
    y0 = array([phi0, phi1])

    # stepsize
    h = double(tEnd) / N
    # memory allocation
    y = zeros((N + 1, 2))
    y[0, :] = y0
    t = zeros(N + 1)

    #########################################
    #                                       #
    # Implement 1 step Stoermer Verlet here #
    #                                       #
    #########################################


    if flag:
        return y[-1][0]
    else:
        return (y, t)


def IntegratePendulumVelocityVerlet(phi0,phi1, N, tEnd=1.8, l=0.6, g=9.81, flag=False):
    """IntegratePendulum solve the mathematical pendulum with two step Stoermer Verlet
    flag == false   -->  y = complete solution phi,phi',t
    flag == true    -->  y = phi(tEnd)
    """
    if shape(phi0) == (1,):
        phi0 = phi0[0]
        phi1 = phi1[0]
    y0 = array([phi0, phi1])

    # stepsize
    h = double(tEnd) / N
    # memory allocation
    y = zeros((N + 1, 2))
    y[0, :] = y0
    t = zeros(N + 1)

    #########################################
    #                                       #
    # Implement velocity Verlet here        #
    #                                       #
    #########################################
    
    if flag:
        return y[-1][0]
    else:
        return (y, t)





if __name__ == '__main__':
    from time import time

    starttime = time()

    tEnd = 1000  # 2*1.8
    N = int(tEnd/0.1)
    l = 0.6
    g = 9.81
    phiT = 5*pi/6.
    phi1 = 0.

    # ODE45
    yODE45, tODE45 = IntegratePendulum(phiT,phi1, tEnd, l, g, False)
    # 2 step Stoermer Verlet
    y_2SV, t_2SV = IntegratePendulumStVerlet2step(phiT,phi1, N, tEnd, l, g, False)
    # 1 step Stoermer Verlet
    y_1SV, t_1SV = IntegratePendulumStVerlet1step(phiT,phi1, N, tEnd, l, g, False)
    # Velocity Verlet
    y_VV, t_VV = IntegratePendulumVelocityVerlet(phiT,phi1, N, tEnd, l, g, False)

    plt.figure()
    plt.plot(yODE45.transpose()[0], yODE45.transpose()[1], 'g-', label='ODE45')
    plt.plot(y_2SV.transpose()[0], y_2SV.transpose()[1], 'b--', label='2SV')
    plt.plot(y_1SV.transpose()[0], y_1SV.transpose()[1], 'r--', label='1SV')
    plt.plot(y_VV.transpose()[0], y_VV.transpose()[1], 'm--', label='VV')

    plt.xlabel(r'$\alpha$')
    plt.ylabel(r'$p$')
    plt.legend()
    plt.axis('equal')
    plt.savefig('solutionsSV.pdf')
    plt.show()

    # Tracking total energy

    # 2 step Stoermer Verlet
    y = y_2SV
    t = t_2SV

    E_kin = 0.5 * (y[:, 1] ** 2)
    E_pot = -g / l * cos(y[:, 0])
    E_pot = E_pot - min(E_pot) + min(E_kin)
    E_tot = E_kin + E_pot

    plt.figure()
    plt.plot(t, E_tot.transpose(), 'b-', label=' 2SV')

    # 1 step Stoermer Verlet
    y = y_1SV
    t = t_1SV

    E_kin = 0.5 * (y[:, 1] ** 2)
    E_pot = -g / l * cos(y[:, 0])
    E_pot = E_pot - min(E_pot) + min(E_kin)
    E_tot = E_kin + E_pot


    plt.plot(t, E_tot.transpose(), 'r-', label='1SV')

    # Velocity Verlet
    y = y_VV
    t = t_VV

    E_kin = 0.5 * (y[:, 1] ** 2)
    E_pot = -g / l * cos(y[:, 0])
    E_pot = E_pot - min(E_pot) + min(E_kin)
    E_tot = E_kin + E_pot

    plt.plot(t, E_tot.transpose(), 'm-', label='VV')

    # ODE45
    y = yODE45
    t = tODE45

    E_kin = 0.5 * (y[:, 1] ** 2)
    E_pot = -g / l * cos(y[:, 0])
    E_pot = E_pot - min(E_pot) + min(E_kin)
    E_tot = E_kin + E_pot

    plt.plot(t, E_tot.transpose(), 'g-', label='ODE45')

    
    plt.xlabel(r'time $t$')
    plt.ylabel('energy')
    plt.legend(loc=2)
    plt.savefig('energies.pdf')
    plt.show()