import numpy as np
import scipy.integrate as sc
import matplotlib.pyplot as plt

g = 9.81 # accélération de la pesanteur
m = 1 # 1 kg
k = 0.01 # coefficient lambda du TD
tau = m/k # m/lambda du TD
v_lim = -m*g/k # vitesse limite du régime permanent portée sur z

def Positif(X, Z):
    """ Renvoie les listes de la partie positive de la trajectoire Z(X) """
    X_pos = []
    Z_pos = []
    for i in range(len(Z)-1):
        if Z[i]>0:
            X_pos.append(X[i])
            Z_pos.append(Z[i])
    return X_pos, Z_pos

def Angles():
    """ Affiche les trajectoires pour différents angles pour le modèle parabolique en rouge et pour le frottement fluide en bleu """
    v0 = 100 # 100 m/s
    plt.close() # ferme la fenêtre courante
    plt.clf() # efface figure courante
    for angle in [10*i for i in range(9)]: # angles entre 0 et 80 °
        alpha = angle # angle initial (°)
        alpha = alpha*2*np.pi/360 # angle initial (rad)
        v0x = v0*np.cos(alpha)
        v0z = v0*np.sin(alpha)

        Xf = 1000 # jusqu'à 1000 m
        N = 1000 # 1000 points
        X = np.linspace(0,Xf,N) # liste des ascisses x
        Z = (v0z-v_lim)/v0x*X - tau*v_lim*np.log(1-X/(tau*v0x)) # trajectoire
        X2 = np.linspace(0,Xf,N) # liste des ascisses x
        Z2 = -g/(2*v0**2*np.cos(alpha)**2)*X**2+np.tan(alpha)*X # parabole
        X, Z = list(X), list(Z) # conversion pour la fonction Positif
        X2, Z2 = list(X2), list(Z2)

        X, Z = Positif(X,Z) # on ne garde que les positions au dessus du sol
        Xpara, Zpara = Positif(X2,Z2) # idem pour la parabole

        plt.axis([0,1000,0,500]) # graphe jusqu'à x = 1000 et y = 500
        plt.title("Variation de l'angle")
        plt.plot(X,Z,"b--")
        plt.plot(Xpara,Zpara,"r-")
        plt.show()

def Vitesses():
    """ Affiche les trajectoires pour différentes vitesses initiales pour le modèle parabolique en rouge et pour le frottement fluide en bleu """
    alpha = 30 # angle initial (°)
    alpha = alpha*2*np.pi/360 # angle initial (rad)
    plt.close() # ferme la fenêtre courante
    plt.clf() # efface figure courante
    for vitesse in [50+10*i for i in range(9)]: # vitesses entre 50 et 130 m/s
        v0 = vitesse
        v0x = v0*np.cos(alpha)
        v0z = v0*np.sin(alpha)

        Xf = 2000 # jusqu'à 1000 m
        N = 1000 # 1000 points
        X = np.linspace(0,Xf,N) # liste des ascisses x
        Z = (v0z-v_lim)/v0x*X - tau*v_lim*np.log(1-X/(tau*v0x)) # trajectoire
        X2 = np.linspace(0,Xf,N) # liste des ascisses x
        Z2 = -g/(2*v0**2*np.cos(alpha)**2)*X**2+np.tan(alpha)*X # parabole
        X, Z = list(X), list(Z) # conversion pour la fonction Positif
        X2, Z2 = list(X2), list(Z2)

        X, Z = Positif(X,Z) # on ne garde que les positions au dessus du sol
        Xpara, Zpara = Positif(X2,Z2) # idem pour la parabole

        plt.axis([0,2000,0,250]) # graphe jusqu'à x = 2000 et y = 250
        plt.title("Variation de la vitesse initiale")
        plt.plot(X,Z,"b--")
        plt.plot(Xpara,Zpara,"r-")
        plt.show()