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


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


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

def IntegratePendulum_A(phi0, N, tEnd=1.8,omega2=16.35, flag=False):
    """IntegratePendulum solve the approximate mathematical pendulum 
    flag == false   -->  y = complete solution phi,phi'
    flag == true    -->  y = phi(tEnd)
    """
    if shape(phi0) == (1,):
        phi0 = phi0[0]
    y0 = array([phi0, 0])

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

    omega = sqrt(omega2)

    for k in range(N):
        t[k + 1] = (k + 1) * h
        # Implement the first method here
        #y[k+1, 0] = 
        #y[k+1, 1] = 
        

    if flag:
        return y[-1][0]
    else:
        return (y, t)
        
def IntegratePendulumEE(phi0, N, tEnd=1.8,omega2=16.35, flag=False):
    """IntegratePendulum solve the mathematical pendulum with EE
    flag == false   -->  y = complete solution phi,phi'
    flag == true    -->  y = phi(tEnd)
    """
    if shape(phi0) == (1,):
        phi0 = phi0[0]
    y0 = array([phi0, 0])

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

    for k in range(N):
        # Implement EE here
        #y[k+1, :] = 
        t[k + 1] = (k + 1) * h

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


def IntegratePendulumIE(phi0, N, tEnd=1.8,omega2=16.35, flag=False):
    """IntegratePendulum solve the mathematical pendulum with IE
    flag == false   -->  y = complete solution phi,phi'
    flag == true    -->  y = phi(tEnd)
    """
    if shape(phi0) == (1,):
        phi0 = phi0[0]
    y0 = array([phi0, 0])

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

    for k in range(N):
        
        #################################
        #                               #
        # Implement IE here #
        #                               #
        #################################
        #y[k+1, :] = y0
        t[k + 1] = (k + 1) * h

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


def IntegratePendulumIM(phi0, N, tEnd=1.8,omega2=16.35, flag=False):
    """IntegratePendulum solve the mathematical pendulum with Implicit Midpoint
    flag == false   -->  y = complete solution phi,phi'
    flag == true    -->  y = phi(tEnd)
    """
    if shape(phi0) == (1,):
        phi0 = phi0[0]
    y0 = array([phi0, 0])

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

    for k in range(int(N)):
        ####################################
        #                                  #
        # Implement IM here #
        #                                  #
        ####################################
        #y[k+1, :] = y0
        t[k + 1] = (k + 1) * h
        
    if flag:
        return y[-1][0]
    else:
        return (y, t)
    
def plotPendelEnergy(y, t, methodlabel, omega2):
    E_kin = 0.5 * (y[:, 1] ** 2)
    E_pot = - omega2 * 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_kin.transpose(), 'b-', label='kinetic energy')
    plt.plot(t, E_pot.transpose(), 'r-', label='potential energy')
    plt.plot(t, E_tot.transpose(), 'g-', label='total energy')
    plt.xlabel(r'time $t$')
    plt.ylabel('energy')
    plt.legend()
    plt.savefig('energie'+methodlabel+'.pdf')
    plt.show()

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

    starttime = time()

    tEnd = 4  # 2*1.8
    N = 500
    l = 0.6
    g = 9.81
    phiT = 1.48551280723
    omega2 = g/l
    # ODE45
    yODE45, tODE45 = IntegratePendulum(phiT, tEnd, omega2, False)
    #
    y_EE, t_EE = IntegratePendulumEE(phiT, N, tEnd, omega2, False)
    #
    y_IE, t_IE = IntegratePendulumIE(phiT, N, tEnd, omega2, False)
    #
    y_IM, t_IM = IntegratePendulumIM(phiT, N, tEnd, omega2, False)
    # approximate model
    y_A, t_A = IntegratePendulum_A(phiT, N*100, tEnd, omega2, False)

    plt.figure()
    plt.plot(yODE45.transpose()[0], yODE45.transpose()[1], 'g-', label='ODE45')
    plt.plot(y_EE.transpose()[0], y_EE.transpose()[1], 'r--', label='EE')
    plt.plot(y_IE.transpose()[0], y_IE.transpose()[1], 'b--', label='IE')
    plt.plot(y_IM.transpose()[0], y_IM.transpose()[1], 'c--', label='IM')
    plt.plot(y_A.transpose()[0], y_A.transpose()[1], 'k--', label='AP')
    #ax = axis
    # plot([ax(1) ax(2)],[0 0],'k-')
    # plot([0 0],[ax(3) ax(4)],'k-')
    plt.xlabel(r'$\alpha$')
    plt.ylabel(r'$p$')
    plt.legend()
    plt.axis('equal')
    plt.savefig('solutions.pdf')
    plt.show()

    # Tracking total energy
    # Explicit Euler
    plotPendelEnergy(y_EE, t_EE, 'EE', omega2)
    # Implicit Euler
    plotPendelEnergy(y_IE, t_IE, 'IE', omega2)
    # Implicit Midpoint
    plotPendelEnergy(y_IM, t_IM, 'IM', omega2)  
    # Approximate equation
    plotPendelEnergy(y_A, t_A, 'AP', omega2)

    print('eror_AP = ', abs(yODE45[-1]-y_A[-1]))
    print('eror_IM = ', abs(yODE45[-1]-y_IM[-1]))
    print('eror_EE = ', abs(yODE45[-1]-y_EE[-1]))
    print('eror_IE = ', abs(yODE45[-1]-y_IE[-1]))