''' Trajectoire d'un projectile avec frottements proportionnels à v^2'''
## Importation des modules
import numpy as np
import matplotlib.pyplot as plt
from math import *
from matplotlib.widgets import Slider


## Constantes
g = -9.81 
m = 1
beta = 1
dt = 1e-3

## Fonctions
def integre(v0, theta): 
    theta = theta/180*pi
    y = 0
    z = 0
    yp = v0*cos(theta)
    zp = v0*sin(theta)
    Lz = [z]
    Ly = [y]
    while z>-1e-6 : 
        yp_new = yp - beta*sqrt(yp**2 + zp**2)*yp*dt
        zp_new = zp +(g- beta*sqrt(yp**2 + zp**2)*zp)*dt
        y_new = y + yp_new*dt
        z_new = z + zp_new*dt
        z = z_new
        y = y_new
        zp = zp_new
        yp = yp_new
        Lz.append(z)
        Ly.append(y)
    return Lz, Ly
    

def dessine():	
    theta = amplitude2.val
    v0 = amplitude1.val
    Lz, Ly = integre(v0, theta)
    l1.set_xdata(Ly)
    l1.set_ydata(Lz)
    ax.set_ylabel(r'z (m)' , fontsize=20)
    ax.set_xlabel(r'y (m)' , fontsize=20)
    plt.draw()

def slider_changed(event):
    dessine()


## En avant pour faire le tracé. 
plt.close('all')

fig = plt.figure(figsize=[8,8])
ax = fig.add_axes([0.12,0.25,0.8,0.7])
plt.axis([0,2,0,3])
l1,l2 =plt.plot([], [], 'k', [], [], 'k')

sld_pos1 = [0.1, 0.02, 0.8, 0.05]    
sld_pos2 = [0.1, 0.07, 0.8, 0.05]


sld_amplitude1 = plt.axes(sld_pos1, facecolor='grey')
amplitude1 = Slider(sld_amplitude1, r'$v0$', 0.1, 20.0, valinit=0.5)
amplitude1.on_changed(slider_changed)

sld_amplitude2 = plt.axes(sld_pos2, facecolor='grey')
amplitude2 = Slider(sld_amplitude2, r'$theta$', 0, 90.0, valinit=45)
amplitude2.on_changed(slider_changed)

dessine()
plt.show()