import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Text3D
from scipy.integrate import odeint
###############################################################################
# Ox vers le nord, Oy vers l'ouest, Oz vertical du lieu
# Système d'équations différentielles x''= F(x,x',y,y',z,z',t)
#                                     y''= G(x,x',y,y',z,z',t)
#                                     z''= H(x,x',y,y',z,z',t)

# Constantes

g = 9.81                      # Champ de pesanteur                          (m.s-2)              
omega = 7.27e-5               # Vitesse angulaire de rotation de la Terre   (rad.s-1)
latitude = 48*2*3.1416/360    # Latitude                                    (rad)
v0 = 300                      # Vitesse initiale                            (m)

# Fonctions intervenant dans le système d'équations différentielles

def F(X,t):
    x,vx,y,vy,z,vz = X
    return 2*omega*np.sin(latitude)*vy
    
def G(X,t):
    x,vx,y,vy,z,vz = X
    return 2*omega*(-np.sin(latitude)*vx+np.cos(latitude)*vz)

def H(X,t):
    x,vx,y,vy,z,vz = X
    return -g-2*omega*np.cos(latitude)*vy

# Définition de la fonction vectorielle telle que X' = f(X,t)

def f(X,t):
    x,vx,y,vy,z,vz = X
    return [vx,F(X,t),vy,G(X,t),vz,H(X,t)]

t0 = 0                           # Instant initial  (s)
tf = 2*v0/g                      # Instant final    (s) 
t = np.linspace(t0,tf,1000)      # Tableau des instants 

condini = [0,0,0,0,0,v0]          # Conditions initiales [x0,vx0,y0,vy0,z0,vz0]

Solution = odeint(f,condini,t)   # Résolution

x = Solution[:,0]                # Tableau des abscisses en fonction du temps
y = Solution[:,2]                # Tableau des ordonnées en fonction du temps
z = Solution[:,4]                # Tableau des cotes en fonction du temps

print("La déviation vers l'ouest vaut :",y[len(t)-1],'m')
################################################################################
# Tracés de x(t),y(t) et z(t)

plt.subplot(1,3,1)
plt.xlabel('t')
plt.ylabel('x(t)')
plt.grid()
plt.plot(t,x)

plt.subplot(1,3,2)
plt.xlabel('t')
plt.ylabel('y(t)')
plt.grid()
plt.plot(t,y)

plt.subplot(1,3,3)
plt.xlabel('t')
plt.ylabel('z(t)')
plt.grid()
plt.plot(t,z)

plt.show()
################################################################################
# Tracés de z(x) et z(y)

plt.subplot(1,2,1)
plt.xlabel('x')
plt.ylabel('z')
plt.grid()
plt.plot(x,z)

plt.subplot(1,2,2)
plt.xlabel('y')
plt.ylabel('z')
plt.grid()
plt.plot(y,z)

plt.show()
################################################################################
# Tracé de la trajectoire en trois dimensions

fig = plt.figure()
axe = Axes3D(fig)

plt.ticklabel_format(axis='x', style='sci', scilimits=(0, 0))
axe.set_xlabel('x')
axe.set_ylabel('y')
axe.set_zlabel('z')

axe.plot(x,y,z)

plt.show()
################################################################################