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

import numpy as np

# Constantes physiques: accélération gravitationnelle, angle d'ouverture, 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.

# Moment cinétique, cos et sin de alpha
l = m * r0**2 * omega0
sinalpha, cosalpha = np.sin(alpha), np.cos(alpha)

def rk4(f, x0, h, N):
    x = np.empty([N+1, 3])  # variables dynamiques
    x[0] = x0  # nb. x[i] = i-ème ligne du tableau x
    for n in range(N):
        k1 = f(x[n])
        k2 = f(x[n] + h/2 * k1)
        k3 = f(x[n] + h/2 * k2)
        k4 = f(x[n] + h * k3)
        x[n+1] = x[n] + h * (k1/6 + k2/3 + k3/3 + k4/6)
    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[0], x[1], x[2] 
    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])
    
x0 = np.array([r0, v0, phi0])
solution = rk4(f, x0, h, N)

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, 0.01)
Y = np.arange(-.12, .12, 0.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()