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

if __name__ == "__main__":
    g = lgc.Const(10)
    m = lgc.Const(1)

    x = lgc.Var("x")
    y = lgc.Var("y")

    T = 0.5 * m * (x.diff() * x.diff() + y.diff() * y.diff())
    V = m * g * y
    L = T - V

    dico = {
        "x" : 0,
        "d_x" : 1,
        "y" : 0,
        "d_y" : 1
    }

    values = L.solve(dico, 10, 0.1)

    # Extraire les données pour x et y
    times = sorted(values.keys())
    positions_x = [values[t]["x"] for t in times]
    positions_y = [values[t]["y"] for t in times]

    # Créer la figure
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_xlim(min(positions_x) - 5, max(positions_x) + 5)
    ax.set_ylim(min(positions_y) - 5, max(positions_y) + 5)
    ax.set_xlabel("X (m)")
    ax.set_ylabel("Y (m)")
    ax.grid(True)
    
    # Point animé
    point, = ax.plot([], [], 'ro', markersize=15, label="Particule")
    # Trajectoire
    line, = ax.plot([], [], 'b-', alpha=0.3, linewidth=2, label="Trajectoire")
    # Texte temps
    time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes, fontsize=12)
    
    ax.legend()

    def animate(frame):
        # Point courant
        point.set_data([positions_x[frame]], [positions_y[frame]])
        # Trajectoire jusqu'au point courant
        line.set_data(positions_x[:frame + 1], positions_y[:frame + 1])
        # Texte
        time_text.set_text(f"t = {times[frame]:.2f}s\nx = {positions_x[frame]:.2f}m, y = {positions_y[frame]:.2f}m")
        return point, line, time_text

    ani = animation.FuncAnimation(fig, animate, frames=len(times), 
                                  interval=50, blit=True, repeat=True)
    
    plt.show()