# -*- coding: utf-8 -*-
"""
Created on Sat Oct 12 09:09:24 2024

@author: arnau
"""

import numpy as np
import matplotlib.pyplot as plt


#%% Constante

g =9.80665
l =0.3

t0 = 0.0
tmax = 8
n = 1000
temps = [k*(tmax-t0)/n for k in range(n)]
theta0 = 0.0
omega0 = 8.0

def f(Y):
        fY = np.zeros_like(Y)
        fY[0] = Y[1]
        fY[1] = -g/l* np.sin(Y[0])
        return fY


def energie(theta,omega):
    return 0.5*omega**2+g/l*(1-np.cos(theta))
    
    
#%% Euler
def pendule(theta0,omega0,t0,tmax,n) :
    """
    Parameters
    ----------
    theta0 : float
        l'angle initial.
    omega0 : float
        vitesse angulaire initiale.
    n : nb de subdivisions
    g_sur_ell : float, optional
        le coeffficient g/ell. The default is 1.

    Returns
    -------
    ndarray sol de shape (T,2) avec 
    sol[:,0] contenant les angles
    sol[:,1] contenant les vitesses angulaires
    """
    h = (tmax-t0)/n
    i_max = int(np.ceil(tmax / h))
    # On créé un tableau qu'on va remplir petit à petit
    thetas = np.zeros((i_max, 2))
    # Conditions initiales
    thetas[0, 0] = theta0
    thetas[0, 1] = omega0
    # Iteration
    for i in range(0, i_max - 1):
        # On applique la relation d'Euler vectorielle
        thetas[i + 1] = thetas[i] + h * f(thetas[i])
    return thetas

#%% Méthode de Heun

def pendule_Heun(theta0,omega0,t0,tmax,n) :
    """
    Parameters
    ----------
    theta0 : float
        l'angle initial.
    omega0 : float
        vitesse angulaire initiale.
    n : nb de subdivisions
    g_sur_ell : float, optional
        le coeffficient g/ell. The default is 1.

    Returns
    -------
    ndarray sol de shape (T,2) avec 
    sol[:,0] contenant les angles
    sol[:,1] contenant les vitesses angulaires
    """
    h = (tmax-t0)/n
    i_max = int(np.ceil(tmax / h))
    # On créé un tableau qu'on va remplir petit à petit
    thetas = np.zeros((i_max, 2))
    # Conditions initiales
    thetas[0, 0] = theta0
    thetas[0, 1] = omega0
    # Iteration
    for i in range(0, i_max - 1):
        thetas[i + 1] = thetas[i] + h * f(thetas[i])
        thetas[i+1]=thetas[i] + h/2*(f(thetas[i])+f(thetas[i+1]))
    return thetas

#%% Test et graphe

Sol_Euler = pendule(theta0,omega0,t0,tmax,n)
Sol_Heun = pendule_Heun(theta0,omega0,t0,tmax,n)

Energie_Euler = energie(Sol_Euler[:,0],Sol_Euler[:,1])
Energie_Heun = energie(Sol_Heun[:,0],Sol_Heun[:,1])
plt.plot(temps,Energie_Heun)
plt.show()