# -*- coding: utf-8 -*-
"""
Created on Tue Jan 20 15:55:08 2015

Listing 11.3 : pendule par l'algorithme de Runge Kutta
d'ordre 4

@author: grivet
"""
import numpy as np
import matplotlib.pyplot as plt

def smy(t,y,z):
    return z 

def smz(t,y,z):
    return -np.sin(y)

def rk4(h,ti,yi,zi):
    k1y = smy(ti,yi,zi); 
    k1z = smz(ti,yi,zi);
    k2y = smy(ti+h/2,yi+h*k1y/2,zi+h*k1z/2);
    k2z = smz(ti+h/2,yi+h*k1y/2,zi+h*k1z/2);
    k3y = smy(ti+h/2,yi+h*k2y/2,zi+h*k2z/2);
    k3z = smz(ti+h/2,yi+h*k2y/2,zi+h*k2z/2);
    k4y = smy(ti+h,yi+h*k3y,zi+h*k3z);
    k4z = smz(ti+h,yi+h*k3y,zi+h*k3z);
    yf = yi + (h/6)*(k1y + 2*k2y + 2*k3y + k4y);
    zf = zi + (h/6)*(k1z + 2*k2z + 2*k3z + k4z);
    tf = ti+h;
    return   yf,zf,tf
  
npt = 1200
t = np.zeros((npt,1)) 
y = np.zeros_like(t); z = np.zeros_like(t)
y0 = float(input('position initiale(rd) : '))
z0 = float(input('vitesse initiale(rd/s) : '))
pas = float(input('longueur du pas : '))
t[0] = 0; y[0] = y0; z[0] = z0
for i in range (1,npt):
    y[i],z[i],t[i] = rk4(pas,t[i-1],y[i-1],z[i-1]);

plt.plot(t,y)
plt.title("mouvement d'un pendule simple")
plt.xlabel("temps")
plt.ylabel("angle")
plt.figure(2)
ep = 1-np.cos(y);
ec = 0.5*z*z
et0 = ep[0] + ec[0]
plt.plot(t,ep,label="e_pot")
plt.plot(t,ec,label="e_cin")
plt.plot(t,(ep+ec-et0)*1e5,label="$10^5(e_{tot} - e_{ini})$")
plt.legend()
plt.xlabel("temps")
plt.ylabel("énergies")