import numpy as np
import matplotlib.pyplot as plt

# Parameters
g = 9.81         # aceleración debido a la gravedad (m/s^2)
l1 = 0.2         # longitud del primer péndulo (m)
l2 = 0.5         # longitud del segundo péndulo (m)
m1 = 0.2         # masa del primer péndulo (kg)
m2 = 0.3         # masa del segundo péndulo (kg)

# Condiciones iniciales
th10 = np.pi/2    # ángulo inicial del primer péndulo (rad)
th20 = -np.pi/4    # ángulo inicial del segundo péndulo (rad)
o10 = 0           # velocidad angular inicial del primer péndulo (rad/s)
o20 = 0           # velocidad angular inicial del segundo péndulo (rad/s)

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

# Aceleración angular
def ac(o1, o2, th1, th2, l1, l2, m1, m2, g):
    N1 = -g*(2*m1 + m2)*np.sin(th1) - m2*g*np.sin(th1 - 2*th2) - 2*m2*np.sin(th1 - th2)*(o2**2*l2 + o1**2*l1*np.cos(th1 - th2))
    N2 = 2*np.sin(th1 - th2)*(o1**2*l1*(m1 + m2) + g*(m1 + m2)*np.cos(th1) + o2**2*l2*m2*np.cos(th1 - th2))
    D1 = l1*(2*m1 + m2 - m2*np.cos(2*th1 - 2*th2))
    D2 = l2*(2*m1 + m2 - m2*np.cos(2*th1 - 2*th2))
    ao1 = N1/D1
    ao2 = N2/D2
    return ao1, ao2    

def rk4(l1, l2, m1, m2, th10, th20, o10, o20, dt, t):
    # Arrays to store solutions
    th1 = np.zeros_like(t)
    th2 = np.zeros_like(t)
    o1 = np.zeros_like(t)
    o2 = np.zeros_like(t)

    # Seeds
    th1[0] = th10
    th2[0] = th20
    o1[0] = o10
    o2[0] = o20  
    
    # Solve the equations using the fourth-order Runge-Kutta method
    for i in range(1, len(t)):
        # Calculate intermediate increments for o1 and o2
        ao1, ao2 = ac(o1[i-1], o2[i-1], th1[i-1], th2[i-1], l1, l2, m1, m2, g)
        k1o1 = ao1*dt
        k1o2 = ao2*dt
        k1th1 = o1[i-1]*dt
        k1th2 = o2[i-1]*dt

        ao1, ao2 = ac(o1[i-1]+0.5*k1o1, o2[i-1]+0.5*k1o2, th1[i-1]+0.5*k1th1, th2[i-1]+0.5*k1th2, l1, l2, m1, m2, g)
        k2o1 = ao1*dt
        k2o2 = ao2*dt
        k2th1 = (o1[i-1]+0.5*k1o1)*dt
        k2th2 = (o2[i-1]+0.5*k1o2)*dt

        ao1, ao2 = ac(o1[i-1]+0.5*k2o1, o2[i-1]+0.5*k2o2, th1[i-1]+0.5*k2th1, th2[i-1]+0.5*k2th2, l1, l2, m1, m2, g)
        k3o1 = ao1*dt
        k3o2 = ao2*dt
        k3th1 = (o1[i-1]+0.5*k2o1)*dt
        k3th2 = (o2[i-1]+0.5*k2o2)*dt

        ao1, ao2 = ac(o1[i-1]+k3o1, o2[i-1]+k3o2, th1[i-1]+k3th1, th2[i-1]+k3th2, l1, l2, m1, m2, g)
        k4o1 = ao1*dt
        k4o2 = ao2*dt
        k4th1 = (o1[i-1]+k3o1)*dt
        k4th2 = (o2[i-1]+k3o2)*dt

        # Update velocities and positions
        o1[i] = o1[i-1] + (k1o1 + 2*k2o1 + 2*k3o1 + k4o1)/6
        o2[i] = o2[i-1] + (k1o2 + 2*k2o2 + 2*k3o2 + k4o2)/6
        th1[i] = th1[i-1] + (k1th1 + 2*k2th1 + 2*k3th1 + k4th1)/6
        th2[i] = th2[i-1] + (k1th2 + 2*k2th2 + 2*k3th2 + k4th2)/6

    # Return the additional results
    return th1, th2, o1, o2, t

# Solve the equations using the fourth-order Runge-Kutta method
th1, th2, o1, o2, t = rk4(l1, l2, m1, m2, th10, th20, o10, o20, dt, t)

# Parámetros para el diagrama de bifurcación
l2_values = np.linspace(0.1, 1.0, 20)  # Valores de longitud del segundo péndulo
max_angles = []

for l2 in l2_values:
    # Resolver el sistema para cada valor de l2
    th1, th2, o1, o2, t = rk4(l1, l2, m1, m2, th10, th20, o10, o20, dt, t)
    
    # Calcular el ángulo máximo del primer péndulo para cada valor de l2
    max_angle = np.max(np.abs(th1))
    max_angles.append(max_angle)

# Graficar el diagrama de bifurcación
plt.figure(figsize=(12, 6))
plt.plot(l2_values, max_angles, 'o-', label='Ángulo máximo del primer péndulo')
plt.xlabel('Longitud del segundo péndulo $l2$ (m)', fontsize=12)
plt.ylabel('Ángulo máximo (rad)', fontsize=12)
plt.title('Diagrama de Bifurcación', fontsize=14)
plt.grid(True)
plt.legend()
plt.show()