import numpy as np
import matplotlib.pyplot as plt

G = 9.81
MASSE = 100
GAMMA1 = 1
GAMMA2 = 2
VITESSE_CRITIQUE = 50
V_INIT = 200
Y_INIT = 0

# paramètres de discrétisation
#T_FINAL = 10
#N_ETAPES = 100
#DT = T_FINAL / N_ETAPES
DT = 0.1
PRECISION = 0.00001

def frottement(vitesse):
    """Force de frottement."""
    transition = 1 - np.exp(-vitesse/VITESSE_CRITIQUE)
    frot = -GAMMA1*vitesse - GAMMA2*transition*vitesse**2
    return frot


def frotprime(v):
    """Dérivée de la fonction de frottement."""
    tr = 1 - np.exp(-v/VITESSE_CRITIQUE)
    trprime = np.exp(-v/VITESSE_CRITIQUE) / VITESSE_CRITIQUE
    frotp = -GAMMA1 - GAMMA2*(2*tr*v + trprime*v**2)
    return frotp

def newton1(x, vn):
    """Fonction "f(x)" de la méthode de Newton."""
    return x - vn + G*DT - DT/(2*MASSE) * (
        frottement(x) + frottement(vn))


def newton2(x, vn):
    """Fonction "f'(x)" de la méthode de Newton."""
    return 1 - DT/(2*MASSE) * frotprime(x)


def newtonRK(vn):
    """Faire une étape de RK grâce à la méthode de Newton."""
    x_actuel = vn
    while True:
        x_prec = x_actuel
        x_actuel = x_prec - newton1(x_prec, vn) / newton2(x_prec, vn)
        if abs(x_prec - x_actuel) < PRECISION:
            return x_actuel

def evolution():
    """Simuler la chute libre avec discrétisation RK."""
    temps = [0]
    vitesses = [V_INIT]
    altitudes = [Y_INIT]
    etape = 0
    while True:
        temps.append(temps[etape] + DT)
        vitesses.append(newtonRK(vitesses[etape]))
        trapeze = (vitesses[etape] + vitesses[etape+1])/2 * DT
        altitudes.append(altitudes[etape] + trapeze)
        if altitudes[etape+1] < 0:
            break
        etape += 1
    return temps, altitudes, vitesses


ts, ys, vs = evolution()
fig, axes = plt.subplots(2, 1)
axes[0].plot(ts, vs, "g")
axes[1].plot(ts, ys, "b")
axes[0].set_ylabel("vitesse")
axes[1].set_xlabel("temps")
axes[1].set_ylabel("altitude")
plt.show()