from vpython import *

# -------------------------
# Parámetros del sistema
# -------------------------
g = 9.81         # m/s^2
m = 1          # kg  (masa del bob)
l = 2         # m   (longitud del péndulo)
R = 1         # m   (radio del disco)
omega = 1.72      # rad/s (velocidad angular del disco)

# Condiciones iniciales
phi = 180*pi/180      # rad
phidot = 0.0         # rad/s

# Integración temporal
t = 0.0
dt = 1e-3            # prueba 1e-3; si quieres más precisión baja a 1e-4 (será más lento)
tmax = 60
fps = 300            # velocidad de animación (iteraciones por segundo)

# -------------------------
# Escena y objetos
# -------------------------
scene = canvas(title="Péndulo en el borde de un disco rotante",
               width=1000, height=600, background=color.black)

# “Disco” (solo como guía visual)
disk = ring(pos=vector(0, 0, 0), axis=vector(0, 0, 1),
            radius=R, thickness=0.02*R, color=color.gray(0.5))

pivot = sphere(radius=0.03*R, color=color.white)
bob   = sphere(radius=0.06*R, color=color.yellow, make_trail=True, trail_radius=0.01*R)
rod   = cylinder(radius=0.02*R, color=color.white)

# -------------------------
# Gráfica de energía
# Nota: el sistema es no autónomo -> E(t) NO es constante
# -------------------------
g1 = graph(title="Energía (sistema no autónomo)", xtitle="t [s]", ytitle="E [J]")
fE = gcurve(color=color.blue,  label="E = K + U")
fK = gcurve(color=color.red,   label="K")
fU = gcurve(color=color.green, label="U")

# -------------------------
# Dinámica: phi'' = (R ω^2 cos(phi-ωt) - g sin(phi))/l
# -------------------------
def accel(phi_val, t_val):
    return (R*omega**2*cos(phi_val - omega*t_val) - g*sin(phi_val)) / l

def rk4_step(phi_val, v_val, t_val, dt_val):
    # Sistema: phi' = v ; v' = accel(phi,t)
    k1_phi = v_val
    k1_v   = accel(phi_val, t_val)

    k2_phi = v_val + 0.5*dt_val*k1_v
    k2_v   = accel(phi_val + 0.5*dt_val*k1_phi, t_val + 0.5*dt_val)

    k3_phi = v_val + 0.5*dt_val*k2_v
    k3_v   = accel(phi_val + 0.5*dt_val*k2_phi, t_val + 0.5*dt_val)

    k4_phi = v_val + dt_val*k3_v
    k4_v   = accel(phi_val + dt_val*k3_phi, t_val + dt_val)

    phi_new = phi_val + (dt_val/6.0)*(k1_phi + 2*k2_phi + 2*k3_phi + k4_phi)
    v_new   = v_val   + (dt_val/6.0)*(k1_v   + 2*k2_v   + 2*k3_v   + k4_v)

    return phi_new, v_new

# -------------------------
# Bucle principal
# -------------------------
while t < tmax:
    rate(fps)

    # Integrar (RK4)
    phi, phidot = rk4_step(phi, phidot, t, dt)
    t += dt

    # (Opcional) envolver el ángulo para visualización
    phi = (phi + pi) % (2*pi) - pi

    # Cinemática: pivote en el borde del disco (moviéndose en círculo)
    pivot_pos = vector(R*cos(omega*t), R*sin(omega*t), 0)
    bob_pos   = pivot_pos + vector(l*sin(phi), -l*cos(phi), 0)

    pivot.pos = pivot_pos
    bob.pos   = bob_pos
    rod.pos   = pivot_pos
    rod.axis  = bob_pos - pivot_pos

    # Velocidad del bob en el marco inercial para K
    pivot_v = vector(-R*omega*sin(omega*t), R*omega*cos(omega*t), 0)
    rel_v   = vector(l*cos(phi)*phidot,     l*sin(phi)*phidot,     0)
    v_bob   = pivot_v + rel_v

    K = 0.5*m*mag2(v_bob)
    U = m*g*bob.pos.y
    E = K + U

    fK.plot(t, K)
    fU.plot(t, U)
    fE.plot(t, E)