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

@author: arnau
"""

import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as integr


#%% Constante

g =9.80665
l =0.3

t0 = 0.0
T = 8
n = 1000
theta0 = 0.0
omega0 = 8.0

#%% odeint


def F(X,t):
    return np.array([X[1],-g/l*np.sin(X[0])])

Abs = np.linspace(t0,T,1000)
X = integr.odeint(F,np.array([theta0,omega0]),Abs)
plt.plot(Abs,X[:,0],label="theta avec odeint")
plt.legend()
plt.show()

#%% Euler

def f(Y):
        fY = np.zeros_like(Y)
        fY[0] = Y[1]
        fY[1] = -g/l* np.sin(Y[0])
        return fY
    
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

temps = [k*(T-t0)/n for k in range(n)]
Sol_Euler = pendule(theta0,omega0,t0,T,n)
plt.plot(temps,Sol_Euler[:,0],label='theta Euler')
plt.legend()
plt.show()

#%% Energie


def energie(theta,omega):
    return 0.5*omega**2+g/l*(1-np.cos(theta))

Energie_Euler = energie(Sol_Euler[:,0],Sol_Euler[:,1])
Energie_scipy = energie(X[:,0],X[:,1])

plt.plot(temps, Energie_Euler, label='energie Euler')
plt.plot(Abs,Energie_scipy, label='energie avec odeint')
plt.legend()
plt.show()