#!/usr/bin/python3
#
# Équations de mouvement pour une bille sur un étonnoir

import numpy as np

# Constantes physiques: accélération gravitationnelle, pente du cône, masse de la bille
g = 9.81
alpha = np.pi/4
m = 1.E-3

# Constantes numériques: pas d'incrément, nombre de pas à calculer
h = 1.E-3
N = 5000

# Conditions initiales: rayon, vitesse radiale, angle, vitesse angulaire à t = 0
r0 = .1
v0 = .2
phi0 = 0.0
omega0 = 4.

# Constantes dépendantes: moment cinétique, sin et cos de alpha
l = m * r0**2 * omega0
sinalpha, cosalpha = np.sin(alpha), np.cos(alpha)

def euler(f, x0, h, N):
    x = np.empty([N+1, 3])  # un tableau pour stocker les variables dynamiques
    x[0] = x0   # nb. x[i] = i-ème ligne du tableau = vecteur à 3 composantes!
    for n in range(N):
        x[n+1] = x[n] + h * f(x[n])
    return x

# La fonction vectorielle qui représente le membre de droite des e.d.m.
# x = un vecteur à 3 composantes contenant r, v et phi au temps t
def f(x):
    r, v, phi = x 
    fr = v
    fv = l**2 * sinalpha**2 / (m**2 * r**3) - g * sinalpha * cosalpha
    fphi = l / (m * r**2)
    return np.array([fr, fv, fphi])

# Réaliser les conditions initiales et résoudre les e.d.m.    
x0 = np.array([r0, v0, phi0])
solution = euler(f, x0, h, N)

# Affichage
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d
plt.rc('text', usetex=True)
plt.rc('font', family='serif')
plt.tick_params(labelsize=16)
ax = plt.figure(dpi=200).add_subplot(projection='3d')
t = np.linspace(0, h * N, N + 1)
r = solution[:, 0]
phi = solution[:, 2]
z = cosalpha/sinalpha * r
x = r * np.cos(phi)
y = r * np.sin(phi)
X = np.arange(-.12, .12, .01)
Y = np.arange(-.12, .12, .01)
X, Y = np.meshgrid(X, Y)
R = np.sqrt(X**2 + Y**2)
Z = cosalpha/sinalpha * R
surf = ax.plot_surface(X, Y, Z, linewidth=0, antialiased=False, alpha=0.2)
ax.plot(x, y, z)

plt.figure(dpi=200)
plt.xlabel('$t$',fontsize=16)
plt.ylabel('$r(t)$' ,fontsize=16) 
plt.plot(t, r)
plt.show()


    

