import numpy as np
import matplotlib
matplotlib.use("Agg")   # backend sin display (necesario en mycompile)
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from collections import deque

# =====================================================
# PARÁMETROS DEL SISTEMA
# =====================================================
g  = 9.81
m  = 1.0
k  = 40.0
l0 = 1.0

Y0 = np.array([60*np.pi/180, 0.0, l0, 0.0])

tf_sim = 20.0
video_seconds = 20
fps = 30
n_frames = video_seconds * fps

h = 1e-3
N = int(tf_sim / h) + 1

trail_len = 60

# =====================================================
# ECUACIONES DE MOVIMIENTO
# =====================================================
def f(t, Y):
    theta, theta_dot, l, l_dot = Y
    theta_dd = -(g*np.sin(theta) + 2*l_dot*theta_dot) / l
    l_dd = (k/m)*(l0 - l) + g*np.cos(theta) + l*theta_dot**2
    return np.array([theta_dot, theta_dd, l_dot, l_dd])

def rk4_step(t, Y, h):
    k1 = f(t, Y)
    k2 = f(t + h/2, Y + h/2 * k1)
    k3 = f(t + h/2, Y + h/2 * k2)
    k4 = f(t + h, Y + h * k3)
    return Y + (h/6) * (k1 + 2*k2 + 2*k3 + k4)

# =====================================================
# INTEGRACIÓN
# =====================================================
t = 0.0
Y = Y0.copy()

tiempo = np.empty(N)
theta_sol = np.empty(N)
thetad_sol = np.empty(N)
l_sol = np.empty(N)
ld_sol = np.empty(N)

for i in range(N):
    tiempo[i] = t
    theta_sol[i], thetad_sol[i], l_sol[i], ld_sol[i] = Y
    Y = rk4_step(t, Y, h)
    t += h

x_sol = l_sol * np.sin(theta_sol)
y_sol = -l_sol * np.cos(theta_sol)

vx = ld_sol*np.sin(theta_sol) + l_sol*np.cos(theta_sol)*thetad_sol
vy = -ld_sol*np.cos(theta_sol) + l_sol*np.sin(theta_sol)*thetad_sol
K_sol  = 0.5*m*(vx**2 + vy**2)
U_sol  = m*g*y_sol + 0.5*k*(l_sol - l0)**2
E_sol  = K_sol + U_sol

# =====================================================
# FIGURA
# =====================================================
fig, (ax_scene, ax_energy) = plt.subplots(
    2, 1, figsize=(8, 10), gridspec_kw={"height_ratios": [2.2, 1]})

fig.patch.set_facecolor("black")

lim = (l0 + max(np.abs(l_sol - l0).max(), 0.5)) * 1.3
ax_scene.set_facecolor("black")
ax_scene.set_xlim(-lim, lim)
ax_scene.set_ylim(-lim, lim*0.4)
ax_scene.set_aspect("equal")
ax_scene.axis("off")
ax_scene.set_title("Péndulo simple con longitud variable (péndulo elástico)", color="white", fontsize=12)

ax_scene.plot(0, 0, 'o', color="white", ms=6, zorder=6)

trail_line, = ax_scene.plot([], [], color="yellow", lw=1, alpha=0.5, zorder=3)
rod, = ax_scene.plot([], [], color="#f4a261", lw=1.8, zorder=4)
bob, = ax_scene.plot([], [], 'o', color="yellow", ms=14, zorder=5)

trail_x = deque(maxlen=trail_len)
trail_y = deque(maxlen=trail_len)

def spring_coords(x1, y1, x2, y2, n_coils=12):
    dx = x2 - x1
    dy = y2 - y1
    L = np.hypot(dx, dy) + 1e-9
    t = np.linspace(0, 1, 300)
    s = 0.5 * (1 - np.cos(np.pi * t))
    nx = -dy / L
    ny = dx / L
    amp = 0.08 * np.sqrt(l0 / L)
    amp = np.clip(amp, 0.04, 0.10)
    wave = amp * np.sin(2 * np.pi * n_coils * s)
    xs = x1 + dx * s + nx * wave
    ys = y1 + dy * s + ny * wave
    return xs, ys

ax_energy.set_facecolor("black")
ax_energy.set_xlim(0, tf_sim)
y_all_min = min(K_sol.min(), U_sol.min(), E_sol.min())
y_all_max = max(K_sol.max(), U_sol.max(), E_sol.max())
margin = 0.1 * (y_all_max - y_all_min + 1e-9)
ax_energy.set_ylim(y_all_min - margin, y_all_max + margin)
ax_energy.set_title("Energía (sistema autónomo)", color="white", fontsize=11)
ax_energy.set_xlabel("t [s]", color="white")
ax_energy.set_ylabel("E [J]", color="white")
ax_energy.tick_params(colors="white")
for spine in ax_energy.spines.values():
    spine.set_edgecolor("gray")

fE, = ax_energy.plot([], [], color="blue",  lw=1.5, label="E = K + U")
fK, = ax_energy.plot([], [], color="red",   lw=1.2, label="K")
fU, = ax_energy.plot([], [], color="green", lw=1.2, label="U")
ax_energy.legend(fontsize=8, facecolor="black", labelcolor="white", framealpha=0.6)

# =====================================================
# ANIMACIÓN
# =====================================================
def init():
    trail_line.set_data([], [])
    rod.set_data([], [])
    bob.set_data([], [])
    fE.set_data([], [])
    fK.set_data([], [])
    fU.set_data([], [])
    return trail_line, rod, bob, fE, fK, fU

def update(frame):
    sim_idx = int(frame * (N - 1) / (n_frames - 1))
    sim_idx = min(sim_idx, N - 1)

    xb, yb = x_sol[sim_idx], y_sol[sim_idx]

    xs, ys = spring_coords(0, 0, xb, yb)
    rod.set_data(xs, ys)
    bob.set_data([xb], [yb])

    trail_x.append(xb)
    trail_y.append(yb)
    trail_line.set_data(trail_x, trail_y)

    fE.set_data(tiempo[:sim_idx+1], E_sol[:sim_idx+1])
    fK.set_data(tiempo[:sim_idx+1], K_sol[:sim_idx+1])
    fU.set_data(tiempo[:sim_idx+1], U_sol[:sim_idx+1])

    return trail_line, rod, bob, fE, fK, fU

ani = FuncAnimation(fig, update, frames=n_frames, init_func=init,
                     interval=1000/fps, blit=True)

# =====================================================
# GUARDAR COMO GIF (sin necesitar ffmpeg)
# =====================================================
ani.save("pendulo_elastico.gif", writer="pillow", fps=fps, dpi=80)
print("GIF guardado: pendulo_elastico.gif")