import numpy as np
import matplotlib.pyplot as plt

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

# Ángulos iniciales en grados
th10G = 90  # por ejemplo
th20G = -45  # por ejemplo

# Condiciones iniciales
th10 = np.radians(th10G)    # ángulo inicial del primer péndulo (rad)
th20 = np.radians(th20G)    # á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)

# Convert ángulos a coordenadas cartesianas
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=(7, 6))
plt.plot(x2, y2, color='b')
plt.xlabel('$x (m)$', fontsize=12)
plt.ylabel('$y (m)$', fontsize=12)
plt.grid(True)
# Ajustar la escala para que sea real
plt.axis('equal')
plt.tight_layout()
plt.show()

# Plot of the angles as a function of time
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.plot(t, th1, 'b-', linewidth=2)
plt.xlabel('$t\,(s)$', fontsize=12)
plt.ylabel('$\\theta_1 \, (rad)$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)

# Plot of the angles as a function of time
plt.subplot(1, 2, 2)
plt.plot(t, th2, 'r-', linewidth=2)
plt.xlabel('$t\,(s)$', fontsize=12)
plt.ylabel('$\\theta_2 \, (rad)$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.tight_layout()
plt.show()

# Plot del espacio de fases
plt.figure(figsize=(12, 6))
# Espacio de fases para el primer péndulo
plt.subplot(1, 2, 1)
plt.plot(th1, o1, 'b-', linewidth=2)
plt.xlabel('$\\theta_1 \, (rad)$', fontsize=12)
plt.ylabel('$\\omega_1 \, (rad/s)$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)

# Espacio de fases para el segundo péndulo
plt.subplot(1, 2, 2)
plt.plot(th2, o2, 'r-', linewidth=2)
plt.xlabel('$\\theta_2 \, (rad)$', fontsize=12)
plt.ylabel('$\\omega_2 \, (rad/s)$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.tight_layout()
plt.show()

# Graficar resultados
plt.figure(figsize=(8, 6))
plt.plot(th1,th2,'r-', linewidth=2)
plt.xlabel('$\\theta_1 \, (rad)$', fontsize=12)
plt.ylabel('$\\theta_2 \, (rad)$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.tight_layout()
plt.show()

# Construcción de la sección de Poincaré
# Define el plano de corte para la sección de Poincaré
cut_angle = 0  # por ejemplo, ángulo de corte en theta1 = 0 rad

# Identificar las intersecciones de las trayectorias con el plano de corte
poincare_points = []
for i in range(1, len(th1)):
    if (th1[i-1] < cut_angle and th1[i] >= cut_angle) or (th1[i-1] > cut_angle and th1[i] <= cut_angle):
        poincare_points.append((th2[i], o2[i]))

# Convertir los puntos de Poincaré a un array para facilidad de uso en gráficos
poincare_points = np.array(poincare_points)

# Graficar la sección de Poincaré
plt.figure(figsize=(8, 6))
plt.plot(poincare_points[:, 0], poincare_points[:, 1], 'bo', markersize=2)
plt.xlabel('$\\theta_2$ (rad)', fontsize=12)
plt.ylabel('$\\omega_2$ (rad/s)', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.tight_layout()
plt.show()