# -*- coding: utf-8 -*-
"""
Created on Sat Dec  1 13:15:14 2018

@author: nicol
"""

from math import pi, sqrt
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider

def findpts(r, Ep, Em):
    index_min = np.argmin(Ep)
    np.where(Ep <= Em)

# Paramètres
m1 = 3
m2 = 1
G = 1   ## 6.6742e-11 # SI

# Creation de l'axe des abscisses
r = np.linspace(1E-3, 6, 1E4)
theta = np.linspace(-pi, pi, 1E4)

# Creation de la figure
fig, ax = plt.subplots()
plt.subplot(121)
plt.subplots_adjust(left=0.1, bottom=0.25)
plt.axis([0, r[-1], -2, 2]) # Limites des axes (xmin,xmax,ymin,ymax)
plt.xlabel(r'Distance r (u.a)')
plt.ylabel(r'Energie potentielle effective (u.a)')
plt.grid(True, which='both')

# Creation de la fonction a tracer
µ = m1*m2/(m1+m2)
C = 2

Ep = 1/2*µ*C**2/r**2 - G*m1*m2/r
Em = 0*np.ones(len(Ep))

inf = np.where(Ep <= Em)[0]

P1 = Ep[inf[0]]
P2 = Ep[inf[-1]]
if inf[-1] == len(Ep)-1:
    P2 = -10


# Creation de la trace de la fonction s en fonction de abscisses. C'est un objet qui est sauvegarde dans 'l'
l1, = plt.plot(r, Ep, 'r-', linewidth=2)
l2, = plt.plot(r, [0]*len(r), 'b-', linewidth=1.5)
p1, = plt.plot(r[inf[0]], P1, 'go', markersize=14)
p2, = plt.plot(r[inf[-1]], P2, 'go', markersize=14)



# Creation de la figure 2
plt.subplot(122)
plt.subplots_adjust(left=0.1, bottom=0.25)
plt.axis('equal')
plt.axis([-r[-1], 1.5, -r[-1], r[-1]]) # Limites des axes (xmin,xmax,ymin,ymax)

plt.xlabel(r'Distance r (u.a)')
#plt.ylabel(r'Energie potentielle effective (u.a)')
plt.grid(True, which='both')

conique = 1/(1+1*np.cos(theta))

p0, = plt.plot(0, 0, 'ro', markersize=14)
c, = plt.plot(conique*np.cos(theta), conique*np.sin(theta), 'b-', linewidth=1.5)


# Positionnement des barres de modification
axcolor = 'lightgoldenrodyellow'  # Choix de la couleur

ax_e  = plt.axes([0.12, 0.08, 0.75, 0.03])
s_e = Slider(ax_e, r'e', 0, 1.5, valinit=1)

# Fonction de mise a jour du graphique
def update(val):
    e = s_e.val # On recupere la valeur de la barre

    Em = -min(Ep)*(e**2-1)*np.ones(len(Ep))
    
    inf = np.where(Ep <= Em)[0]
    
    P1 = Ep[inf[0]]
    P2 = Ep[inf[-1]]
    if inf[-1] == len(Ep)-1:
        P2 = -10

    conique = 1/(1+e*np.cos(theta))
    if e > 1:
        t = np.where(abs(theta) <= np.arccos(-1/e))[0]
        


    # On met a jour l'objet
    l2.set_ydata(Em)
    p1.set_data(r[inf[0]], P1)
    p2.set_data(r[inf[-1]], P2)
    
    if e > 1:
        1
        c.set_data(conique[t[0]:t[-1]]*np.cos(theta[t[0]:t[-1]]), conique[t[0]:t[-1]]*np.sin(theta[t[0]:t[-1]]))
    else:
        c.set_data(conique*np.cos(theta), conique*np.sin(theta))
    
    fig.canvas.draw_idle() # On provoque la mise a jour du graphique, qui n'est pas automatique par defaut

s_e.on_changed(update) # lorsque la barre s_lamb est modifiee, on applique la fonction update
        
mng = plt.get_current_fig_manager()
#mng.window.state('zoomed')
plt.show() # On provoque l'affichage a l'ecran