import numpy as np
import matplotlib.pyplot as plt

# Parámetros
g = 9.81         # aceleración debido a la gravedad (m/s^2)
m1 = 0.1         # masa del primer péndulo (kg)
m2 = 0.1         # masa del segundo péndulo (kg)
k1 = 10          # constante del resorte entre los dos péndulos
k2 = 10          # constante del resorte entre el primer péndulo y el soporte

# Ángulos y longitudes iniciales
th10G = 10       # ángulo inicial del primer péndulo en grados
th20G = 30      # ángulo inicial del segundo péndulo en grados
l10 = 1       # longitud inicial del primer péndulo (m)
l20 = 1      # longitud inicial del segundo péndulo (m)

# Convertir ángulos a radianes
th10 = np.radians(th10G)
th20 = np.radians(th20G)

# Velocidades angulares y de longitud iniciales
o10 = 0          # velocidad angular inicial del primer péndulo (rad/s)
o20 = 0          # velocidad angular inicial del segundo péndulo (rad/s)
v10 = 0          # velocidad inicial de cambio de l1 (m/s)
v20 = 0          # velocidad inicial de cambio de l2 (m/s)

# Tiempo
t0 = 0
tf = 5
dt = 0.001       # paso de tiempo
t = np.arange(t0, tf, dt)

# Aceleraciones angulares y de longitud
def ac(th1, th2, o1, o2, l1, l2, v1, v2, m1, m2, g, k1, k2):
    # Aceleración para θ1, θ2, l1 y l2 
    a_th1 = -1/(2*l1*m1) * (4*m1*v1*o1 + 2*m1*g*np.sin(th1) + k2*np.sin(th1 - th2))
    a_th2 = -1/(2*l2*m2) * (4*m2*v2*o2 - k1*np.sin(th1 - th2))
    a_l1 = 1/(2*m1) * (2*l1*m1*o1**2 + 2*m1*g*np.cos(th1) + k2*np.cos(th1 - th2) - k1)
    a_l2 = 1/(2*m1*m2) * (2*l2*m1*m2*o2**2 + k1*m2*np.cos(th1 - th2) - k2*(m1 + m2))
    
    return a_th1, a_th2, a_l1, a_l2

# Método de Runge-Kutta de cuarto orden para resolver el sistema
def rk4(m1, m2, th10, th20, o10, o20, l10, l20, v10, v20, dt, t):
    # Arrays para almacenar los resultados
    th1 = np.zeros_like(t)
    th2 = np.zeros_like(t)
    o1 = np.zeros_like(t)
    o2 = np.zeros_like(t)
    l1 = np.zeros_like(t)
    l2 = np.zeros_like(t)
    v1 = np.zeros_like(t)
    v2 = np.zeros_like(t)

    # Condiciones iniciales
    th1[0] = th10
    th2[0] = th20
    o1[0] = o10
    o2[0] = o20
    l1[0] = l10
    l2[0] = l20
    v1[0] = v10
    v2[0] = v20
    
    # Iteración de Runge-Kutta
    for i in range(1, len(t)):
        # Primer paso de Runge-Kutta
        a_th1, a_th2, a_l1, a_l2 = ac(th1[i-1], th2[i-1], o1[i-1], o2[i-1], l1[i-1], l2[i-1], v1[i-1], v2[i-1], m1, m2, g, k1, k2)
        k1o1, k1o2 = a_th1 * dt, a_th2 * dt
        k1v1, k1v2 = a_l1 * dt, a_l2 * dt
        k1th1, k1th2 = o1[i-1] * dt, o2[i-1] * dt
        k1l1, k1l2 = v1[i-1] * dt, v2[i-1] * dt

        # Segundo paso de Runge-Kutta
        a_th1, a_th2, a_l1, a_l2 = ac(th1[i-1] + 0.5 * k1th1, th2[i-1] + 0.5 * k1th2,
                                      o1[i-1] + 0.5 * k1o1, o2[i-1] + 0.5 * k1o2,
                                      l1[i-1] + 0.5 * k1l1, l2[i-1] + 0.5 * k1l2,
                                      v1[i-1] + 0.5 * k1v1, v2[i-1] + 0.5 * k1v2,
                                      m1, m2, g, k1, k2)
        k2o1, k2o2 = a_th1 * dt, a_th2 * dt
        k2v1, k2v2 = a_l1 * dt, a_l2 * dt
        k2th1, k2th2 = (o1[i-1] + 0.5 * k1o1) * dt, (o2[i-1] + 0.5 * k1o2) * dt
        k2l1, k2l2 = (v1[i-1] + 0.5 * k1v1) * dt, (v2[i-1] + 0.5 * k1v2) * dt

        # Actualización final de posiciones y velocidades
        o1[i] = o1[i-1] + (k1o1 + 2 * k2o1) / 6
        o2[i] = o2[i-1] + (k1o2 + 2 * k2o2) / 6
        v1[i] = v1[i-1] + (k1v1 + 2 * k2v1) / 6
        v2[i] = v2[i-1] + (k1v2 + 2 * k2v2) / 6
        th1[i] = th1[i-1] + (k1th1 + 2 * k2th1) / 6
        th2[i] = th2[i-1] + (k1th2 + 2 * k2th2) / 6
        l1[i] = l1[i-1] + (k1l1 + 2 * k2l1) / 6
        l2[i] = l2[i-1] + (k1l2 + 2 * k2l2) / 6
    
    return th1, th2, o1, o2, l1, l2, v1, v2, t

# Resolver usando el método de Runge-Kutta
th1, th2, o1, o2, l1, l2, v1, v2, t = rk4(m1, m2, th10, th20, o10, o20, l10, l20, v10, v20, dt, t)

# Convertir ángulos y longitudes a coordenadas cartesianas para la trayectoria
x1 = l1 * np.sin(th1)
y1 = -l1 * np.cos(th1)
x2 = x1 + l2 * np.sin(th2)
y2 = y1 - l2 * np.cos(th2)

# Graficar la trayectoria en el plano xy
plt.figure(figsize=(8, 6))
plt.plot(x2, y2, color='b')
plt.xlabel('$x \, (m)$', fontsize=12)
plt.ylabel('$y \, (m)$', fontsize=12)
plt.grid(True)
plt.axis('equal')
plt.tight_layout()
plt.show()