#!/usr/bin/python3

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint


def euler(fct, X0, T):
  """
  résolution d'équations différentielles
  fct : fonction definissant l'equation différentielle
  X0  : tableau des condition initiales (1 élément si ED 1er ordre, 2 éléments si 2ième ordre
  T   : tableau des instants à calculer
  return : tableau des angles et vitesses angulaires pour tous les instants T
  """
  X = np.zeros( (n+1, len(X0)) )  # tableau à 2 dimensions: autant de lignes qu'en a X0 et n+1 colonnes pour les n+1 instants
                                  # pour le pendule simple : X[0,k] = theta(tk) et X[1,k]=dtheta/dt(tk)
  X[0,:] = X0                     # la première colonne contient les conditions initiales : theta0 et dtheta/dt(0)
  for i in range(0, n):
    h = T[i+1]-T[i]                 # ici h est constant mais on pourrait l'envisager non constant
    # on calcule le point suivant (colonne i+1) à partir du point précédent (colonne i)
    X[i+1,:] = X[i,:]+h*fct(X[i,:], T[i])
  return X                        # la fonction renvoie le tableau X a 2 dimensions  et  le tableau des temps


def f_pendule(Xt, t):      # Xt[0] correspond à theta(t), Xt[1] correspond à dtheta/dt(t)
  """
  calcul de la dérivée de Xt
  Xt : tableau contenant l'angle et et la vitesse angulaire à l'instant t
  return : tableau contenant vitesse angulaire et accélération angulaire (dérivée de Xt)
  """
  d1 = Xt[1]                 # dérivée de theta(t)
  d2 = -np.sin(Xt[0])        # dérivée seconde donnée par l'ED du second ordre
  return np.array([d1, d2])  # renvoie


############## définition des paramètres généraux
m = 0.1             # kg
l = 0.1             # metres
w0 = np.sqrt(1/m)   # pulsation propre
t0 = 0
tf = 8*2*np.pi/w0   # choix de tf pour représenter 8 périodes pour les petites oscillations
n = 10000           # nombre de pas de calculs

############## résolution et tracé pour une valeur de theta(0), utilisation de la fonction 'euler'
# résolution
theta0 = 2          # rad -> à choisir
dtheta0 = 0         # rad/s -> à choisir
T = np.linspace(t0, tf, n+1)                     # tableau des temps à 1 dimension (n+1 instants de calculs) répartis régulièrement entre t0 et tf
X0 = np.array([theta0, dtheta0])                 # conditions initiales (t=t0)
Xsolution = euler(f_pendule, X0, T)              # résolution par la méthode d'Euler : solution
THETA = Xsolution[:,0]                           # extraction du tableau des angles calculés
DTHETA = Xsolution[:,1]                          # extraction du tableau des vitesses angulaires calculées
# préparation du tracé
fig, ax1 = plt.subplots()
# tracé de l'évolution de theta
ax1.plot(T, THETA, label=r'avec $\theta$(0)=' + str(theta0))
# détails de présentation du graphique
ax1.set_xlabel('t en seconde')
ax1.set_ylabel(r'$\theta$(t) en rad')
ax1.legend(loc='best')
ax1.grid()
plt.title(r"$\theta$(t) pour le pendule avec méthode d'Euler")
plt.show()


############## résolution et tracé pour plusieurs valeurs de theta(0), utilisation de la fonction 'odeint'
# préparation du tracé
fig, ax1 = plt.subplots()
# boucle sur chaque theta(0)
all_theta0 = [.1, .3, 1, 2, 2.8]          # rad -> à choisir
dtheta0 = 0         # rad/s -> à choisir
for theta0 in all_theta0:
  # résolution pour chaque theta(0)
  X0 = np.array([theta0, dtheta0])                 # conditions initiales (t=t0)
  Xsolution = odeint(f_pendule, X0, T)             # résolution par la méthode d'Euler : solution, tableau des temps
  THETA = Xsolution[:,0]                           # extraction du tableau des angles calculés
  DTHETA = Xsolution[:,1]                          # extraction du tableau des vitesses angulaires calculées
  # tracé de l'évolution de theta
  ax1.plot(T, THETA, label=r'avec $\theta$(0)=' + str(theta0))
# détails de présentation du graphique
ax1.set_xlabel('t en seconde')
ax1.set_ylabel(r'$\theta$(t) en rad')
ax1.legend(loc='best')
ax1.grid()
plt.title(r"$\theta$(t) pour le pendule avec méthode odint")
plt.show()