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

g = lgc.Const(-10)
m1 = lgc.Const(1)
m2 = lgc.Const(1)
r10 = lgc.Const(1)
r20 = lgc.Const(1)
k1 = lgc.Const(50)
k2 = lgc.Const(50)

theta1 = lgc.Var("theta1")
theta2 = lgc.Var("theta2")
r1 = lgc.Var("r1")
r2 = lgc.Var("r2")

x1 = r1 * lgc.cos(theta1)
y1 = r1 * lgc.sin(theta1)

x2 = x1 + r2 * lgc.cos(theta2)
y2 = y1 + r2 * lgc.sin(theta2)

T = 0.5 * m1 * (x1.diff()*x1.diff() + y1.diff()*y1.diff())
T += 0.5 * m2 * (x2.diff()*x2.diff() + y2.diff()*y2.diff())

V = m1 * g * x1
V += m2 * g * x2
V += k1 * (r1 - r10) * (r1 - r10)
V += k2 * (r2 - r20) * (r2 - r20)

L = T - V

init = {
    "theta1" : [- np.pi/4, 0],
    "theta2" : [- np.pi/4, 0],
    "r1" : [1, 0],
    "r2" : [1, 0],
}

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

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

x1_positions = [p["r1"] * np.sin(p["theta1"]) for p in positions]
y1_positions = [p["r1"] * np.cos(p["theta1"]) for p in positions]

x2_positions = [x1_positions[i] + p["r2"] * np.sin(p["theta2"]) for i, p in enumerate(positions)]
y2_positions = [y1_positions[i] + p["r2"] * np.cos(p["theta2"]) for i, p in enumerate(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]

# 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 (m)")
ax.set_ylabel("Y (m)")
ax.grid(True)
ax.invert_yaxis()

line1, = ax.plot([], [], 'o-', lw=2, markersize=8, label='Pendule 1')
line2, = ax.plot([], [], 'o-', lw=2, markersize=8, label='Pendule 2')
trace, = ax.plot([], [], 'r-', alpha=0.3, lw=1, label='Trace')

def animate(frame):
    line1.set_data([0, x1_display[frame]], [0, y1_display[frame]])
    line2.set_data([x1_display[frame], x2_display[frame]], [y1_display[frame], y2_display[frame]])
    trace.set_data(x2_display[:frame+1], y2_display[:frame+1])
    return line1, line2, trace

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

plt.show()
plt.show()