import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

g = 9.81
omega = 7.27e-5

#################################################################################
#paramètres modifiables
#################################################################################

latitude_degre = 49
date_max = 7 # temps de chute en secondes
nbe_dates = 10000 # le nombre de points de tracé

#conditions initiales

x0 = 0
y0 = 0
z0 = 209
vx0 = 0
vy0 = 0
vz0 = 0

#################################################################################
#fin des paramètres modifiables
#################################################################################

def equations(X,t):
    #X est (x,y,z,vx,vy,vz)
    #equations renvoie dX/dt
    x = X[0]
    y = X[1]
    z = X[2]
    vx = X[3]
    vy = X[4]
    vz = X[5]
    ax = 2*vy*omega*np.sin(latitude_degre*np.pi/180)
    ay = -2*vx*omega*np.sin(latitude_degre*np.pi/180)-2*vz*omega*np.cos(latitude_degre*np.pi/180)
    az = -g+2*vy*omega*np.cos(latitude_degre*np.pi/180)

    return np.array([vx,vy,vz,ax,ay,az])


# Résolution par odeint

conditions_initiales = [x0,y0,z0,vx0,vy0,vz0]
dates = np.linspace(0,date_max,nbe_dates)
res = odeint(equations,conditions_initiales,dates)

# On récupère x,y,z ainsi que vx,vy,vz

x = res[:,0]
y = res[:,1]
z = res[:,2]
vx = res[:,3]
vy = res[:,4]
vz = res[:,5]

# Fonction permettant de retrouver la déviation au niveau du sol (en z=0)

def dichotomie(T,a,b):
    ''' Recherche par dichotomie de l'indice d'une valeur nulle dans un tableau
     T de flottants ordonnés par valeurs croissantes ou décroissantes '''
    assert T[a]*T[b] < 0
    while abs(a - b) > 1 :
        m = int((a + b)/2.)
        if T[m] == 0.:
            return m
        elif T[a]*T[m] > 0:
            a = m
        else:
            b = m
    return m

# Tracé des lois horaires

plt.close()
plt.figure()
# Changer l'ordonnée x ou y ou z pour représenter l'évolution des trois coordonnées

plt.plot(dates,z)
plt.xlabel("t(s)")
plt.ylabel("z(m)")
plt.title("Déviation vers l'Est'")
plt.grid()
plt.show()