import pylab as plt
import numpy as np
import matplotlib as mpl
import scipy.integrate
from matplotlib.widgets import Slider, Button

#mpl.rcParams['text.usetex'] = False #True pour mettre du beau LateX mais ça fait laguer ou False c'est moche mais ça lag pas
mpl.rcParams['axes.titlesize'] = 15
mpl.rcParams['axes.labelsize'] = 30
mpl.rcParams['lines.linewidth'] = 3
mpl.rcParams['lines.markersize'] = 8
mpl.rcParams['xtick.labelsize'] = 20
mpl.rcParams['ytick.labelsize'] = 20
mpl.rcParams['legend.fontsize'] = 25

#Résolution des équa diff régissant deux pendules de torsion couplés où l'on peut faire varier les conditions initiales

##Le gros du calcul

# The parametrized function to be plotted
def to_plot(theta_1_ini,theta_2_ini,t_final) : 
    #Paramètres physiques qui correspondent normalement à ceux du pendule de TP
    J0 = 0.07 
    M = 500e-3
    L = 50e-2
    C = 0.34
    
    g = 9.81
    
    #Pulsations propres du système couplé
    omega_0 = np.sqrt((M*g*L)/(J0+M*L**2))
    omega_c = np.sqrt(C/(J0+M*L**2))
    
    #Fonction à intégrer
    def f(var , t , omega_0 , omega_c ) :
        theta_1 , theta_1_p , theta_2 , theta_2_p = var[0] , var[1] , var[2] , var[3]
        theta_1_pp = - omega_0**2 * np.sin(theta_1) + omega_c**2 * (theta_2 - theta_1)
        theta_2_pp = - omega_0**2 * np.sin(theta_2) - omega_c**2 * (theta_2 - theta_1)
        return np.array([ theta_1_p , theta_1_pp , theta_2_p , theta_2_pp ])
       
    #Conditions initiales
    var_ini = np.array([theta_1_ini,0,theta_2_ini,0])
    
    #Temps
    t = np.linspace(0,t_final,5000)
    
    #Calcul de la solution avec odeint
    sol = scipy.integrate.odeint(f , var_ini, t, args=(omega_0,omega_c ))
    
    return t,sol[:,0],sol[:,2]
     
##Le plot avec sliders

# Define initial parameters
init_theta_1_ini = np.pi/4
init_theta_2_ini = np.pi/4
init_t_final = 10

# Create the figure and the line that we will manipulate
fig, ax = plt.subplots()
t,theta_1,theta_2 = to_plot(init_theta_1_ini,init_theta_2_ini,init_t_final)
line, = plt.plot(t,theta_1)
line2, = plt.plot(t,theta_2)

ax.set_xlabel('Temps '+r'$t$')
ax.set_ylabel('Angles '+r'$\theta_1$'+' et '+r'$\theta_2$')
ax.set_xlim(0,init_t_final)
ax.set_ylim(-np.pi,np.pi)

# Adjust the main plot to make room for the sliders
#plt.subplots_adjust(left=0.25, bottom=0.25)
plt.subplots_adjust(left=0.1,bottom=0.25,top=0.95)

# Make horizontal sliders to control the parameters.
ax_theta_1_ini = plt.axes([0.1, 0.07, 0.8, 0.03])
theta_1_ini_slider = Slider(
    ax=ax_theta_1_ini,
    label=r'$\theta_{1,0}$',
    valmin=-np.pi,
    valmax=np.pi,
    valinit=init_theta_1_ini,
)
theta_1_ini_slider.label.set_size(20)
theta_1_ini_slider.valtext.set_fontsize(20)

ax_theta_2_ini = plt.axes([0.1, 0.1, 0.8, 0.03])
theta_2_ini_slider = Slider(
    ax=ax_theta_2_ini,
    label=r'$\theta_{2,0}$',
    valmin=-np.pi,
    valmax=np.pi,
    valinit=init_theta_2_ini,
)
theta_2_ini_slider.label.set_size(20)
theta_2_ini_slider.valtext.set_fontsize(20)

ax_t_final = plt.axes([0.1, 0.13, 0.8, 0.03])
t_final_slider = Slider(
    ax=ax_t_final,
    label=r'$t_f$',
    valmin=0.1,
    valmax=100,
    valinit=init_t_final,
)
t_final_slider.label.set_size(20)
t_final_slider.valtext.set_fontsize(20)

# The function to be called anytime a slider's value changes
def update(val):
    t,theta_1,theta_2 = to_plot(theta_1_ini_slider.val, theta_2_ini_slider.val , t_final_slider.val)
    line.set_ydata(theta_1)
    line2.set_ydata(theta_2)
    line.set_xdata(t)
    line2.set_xdata(t)
    ax.set_xlim(0,t_final_slider.val)
    fig.canvas.draw_idle()


# register the update function with each slider
theta_1_ini_slider.on_changed(update)
theta_2_ini_slider.on_changed(update)
t_final_slider.on_changed(update)

# Create a `matplotlib.widgets.Button` to reset the sliders to initial values.
resetax = plt.axes([0.8, 0.025, 0.1, 0.04])
button = Button(resetax, 'Reset', hovercolor='0.975')

def reset(event):
    theta_1_ini_slider.reset()
    theta_2_ini_slider.reset()
    t_final_slider.reset()
button.on_clicked(reset)

plt.show()