import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import legrandchien as lgc
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np

x1 = lgc.Var()
y1 = lgc.Var()
x2 = lgc.Var()
y2 = lgc.Var()
x3 = lgc.Var()
y3 = lgc.Var()

m1 = lgc.Const(1)
m2 = lgc.Const(1)
m3 = lgc.Const(1)
G = 1

T = 0.5 * m1 * (x1.diff()**2 + y1.diff()**2)
T += 0.5 * m2 * (x2.diff()**2 + y2.diff()**2)
T += 0.5 * m3 * (x3.diff()**2 + y3.diff()**2)

r12 = ((x1 - x2)**2 + (y1 - y2)**2)**0.5
r13 = ((x1 - x3)**2 + (y1 - y3)**2)**0.5
r23 = ((x2 - x3)**2 + (y2 - y3)**2)**0.5

V = -G * m1 * m2 / r12
V += -G * m1 * m3 / r13
V += -G * m2 * m3 / r23

L = T - V

init = {
    x1: [1, -0.5],
    y1: [1, 0.5],
    x2: [-1, -0.25],
    y2: [-0.5, 0.5],
    x3: [0, 0.5],
    y3: [0, -1],
}

dt = 0.01
data = L.solve(init, 10, dt)

# Extraire les positions
times = sorted(data.keys())
positions = [data[t] for t in times]

x1_positions = [p["x1"] for p in positions]
y1_positions = [p["y1"] for p in positions]
x2_positions = [p["x2"] for p in positions]
y2_positions = [p["y2"] for p in positions]
x3_positions = [p["x3"] for p in positions]
y3_positions = [p["y3"] for p in positions]

fps = 25
frame_skip = max(1, int(1 / (fps * dt)))

display_indices = list(range(0, len(x1_positions), frame_skip))
x1_display = [x1_positions[i] for i in display_indices]
y1_display = [y1_positions[i] for i in display_indices]
x2_display = [x2_positions[i] for i in display_indices]
y2_display = [y2_positions[i] for i in display_indices]
x3_display = [x3_positions[i] for i in display_indices]
y3_display = [y3_positions[i] for i in display_indices]

# Créer la figure
fig, ax = plt.subplots(figsize=(8, 8))
ax.set_xlim(-3, 3)
ax.set_ylim(-3, 3)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.grid(True)

point1, = ax.plot([], [], 'o', markersize=10, label='Corps 1')
point2, = ax.plot([], [], 'o', markersize=10, label='Corps 2')
point3, = ax.plot([], [], 'o', markersize=10, label='Corps 3')

def animate(frame):
    point1.set_data([x1_display[frame]], [y1_display[frame]])
    point2.set_data([x2_display[frame]], [y2_display[frame]])
    point3.set_data([x3_display[frame]], [y3_display[frame]])
    return point1, point2, point3

ax.legend()
anim = animation.FuncAnimation(fig, animate, frames=len(x1_display), interval=1000//fps, blit=True, repeat=True)

plt.show()