import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as spi

# ----------------------------
# Paramètres SI
# ----------------------------
m = 3.3e23          # masse Mercure, kg
r0 = 46e9           # distance rp Soleil-Mercure au périhélie, m
v0 = 59e3           # vitesse vp au périhélie, m/s
MS = 1.99e30        # masse du Soleil, kg
G = 6.67e-11        # constante gravitationnelle
alpha = 2e19        # coefficient de perturbation, m²

mu = G * MS

# ----------------------------
# Temps
# ----------------------------
N = 10000
t = np.linspace(0, 365 * 24 * 3600, N)

# ----------------------------
# Équations différentielles
# ----------------------------
def f(Y, t):
    x, y, vx, vy = Y

    # A compléter

    return [vx, vy, ax, ay]

# ----------------------------
# Conditions initiales au périhélie
# ----------------------------
# Mercure est initialement en (r0, 0)
# Sa vitesse est perpendiculaire au rayon : (0, v0)
Y0 = [r0, 0, 0, v0]

# ----------------------------
# Appel à odeint
# ----------------------------
Y = spi.odeint(f, Y0, t)

# Récupération des solutions

x = Y[:, 0]
y = Y[:, 1]
vx = Y[:, 2]
vy = Y[:, 3]

r = np.sqrt(x**2 + y**2)

# ----------------------------
# Q3 : Tracé de r(t)
# ----------------------------
plt.figure(figsize=(8, 5))
plt.plot(t / (24 * 3600), r / 1e9, "r", linewidth=2)
plt.xlabel("t (jours)")
plt.ylabel("r(t) en 10⁹ m")
plt.title("Distance Soleil-Mercure en fonction du temps")
plt.grid(True)
plt.show()

# ----------------------------
# Q6 : Tracé de l'orbite
# ----------------------------
plt.figure(figsize=(6, 6))
plt.plot(x / 1e9, y / 1e9, "b", linewidth=0.5)
plt.scatter(0, 0, color="orange", label="Soleil")
plt.xlabel("x en 10⁹ m")
plt.ylabel("y en 10⁹ m")
plt.title("Trajectoire de Mercure")
plt.axis("equal")
plt.grid(True)
plt.legend()
plt.show()