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

g = lgc.Const(10)
m = lgc.Const(1)
k = lgc.Const(10)
r0 = lgc.Const(1)

theta = lgc.Var("theta")
r = lgc.Var("r")

x = r * lgc.cos(theta)
y = r * lgc.sin(theta)

T = 0.5 * m * (x.diff()*x.diff() + y.diff()*y.diff())
V = m * g * r * lgc.cos(theta) + k * (r - r0) * (r - r0)

L = T - V

dico = {
    "theta" : -3.14/2,
    "d_theta" : 0,
    "r" : 1.2,
    "d_r" : 0,
}

data = L.solve(dico, 10, 0.01)

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

x_positions = [p["r"] * np.sin(p["theta"]) for p in positions]
y_positions = [p["r"] * np.cos(p["theta"]) for p in positions]

# Créer la figure et l'axe
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)

line, = ax.plot([], [], 'o-', lw=2, markersize=8)
trace, = ax.plot([], [], 'r-', alpha=0.3, lw=1)

def animate(frame):
    line.set_data([0, x_positions[frame]], [0, y_positions[frame]])
    trace.set_data(x_positions[:frame+1], y_positions[:frame+1])
    return line, trace

anim = animation.FuncAnimation(fig, animate, frames=len(x_positions), interval=50, blit=True, repeat=True)

plt.show()