import numpy as np
import matplotlib.pyplot as plt
from decimal import Decimal, getcontext

# Establecer precisión a 30 decimales
getcontext().prec = 30

# Parameters
g = 9.81         # acceleration due to gravity (m/s²)
mu = 1.1         # reduced mass 

# Initial conditions
r0 = 1         # initial radial distance (m)
th0 = np.pi/2  # initial angle (np.pi rad)
dr0 = 0        # initial radial velocity (m/s)
dth0 = 0       # initial angular velocity (rad/s)

# Time parameters
t0 = 0
tf = 60
dt = 0.001       # time step
t = np.arange(t0, tf, dt)

# Acceleration functions
def ac(r, th, dr, dth, mu, g):
    d2r = (r*dth**2 + g*(np.cos(th)-mu))/(1 + mu)
    d2th = -(2*dr*dth + g*np.sin(th))/r
    return d2r, d2th

def rk4(r0, th0, dr0, dth0, mu, g, dt, t):
    # Arrays to store solutions
    r = np.zeros_like(t)
    th = np.zeros_like(t)
    dr = np.zeros_like(t)
    dth = np.zeros_like(t)

    # Initial conditions
    r[0] = r0
    th[0] = th0
    dr[0] = dr0
    dth[0] = dth0

    # Solve the equations using the fourth-order Runge-Kutta method
    for i in range(1, len(t)):
        # Calculate intermediate increments
        d2r, d2th = ac(r[i-1], th[i-1], dr[i-1], dth[i-1], mu, g)
        k1dr = d2r*dt
        k1dth = d2th*dt
        k1r = dr[i-1]*dt
        k1th = dth[i-1]*dt

        d2r, d2th = ac(r[i-1]+0.5*k1r, th[i-1]+0.5*k1th, dr[i-1]+0.5*k1dr, dth[i-1]+0.5*k1dth, mu, g)
        k2dr = d2r*dt
        k2dth = d2th*dt
        k2r = (dr[i-1]+0.5*k1dr)*dt
        k2th = (dth[i-1]+0.5*k1dth)*dt
        
        d2r, d2th = ac(r[i-1]+0.5*k2r, th[i-1]+0.5*k2th, dr[i-1]+0.5*k2dr, dth[i-1]+0.5*k2dth, mu, g)
        k3dr = d2r*dt
        k3dth = d2th*dt
        k3r = (dr[i-1]+0.5*k2dr)*dt
        k3th = (dth[i-1]+0.5*k2dth)*dt

        d2r, d2th = ac(r[i-1]+k3r, th[i-1]+k3th, dr[i-1]+k3dr, dth[i-1]+k3dth, mu, g)
        k4dr = d2r*dt
        k4dth = d2th*dt
        k4r = (dr[i-1]+k3dr)*dt
        k4th = (dth[i-1]+k3dth)*dt
        
        # Update values
        dr[i] = dr[i-1] + (k1dr + 2*k2dr + 2*k3dr + k4dr)/6
        dth[i] = dth[i-1] + (k1dth + 2*k2dth + 2*k3dth + k4dth)/6
        r[i] = r[i-1] + (k1r + 2*k2r + 2*k3r + k4r)/6
        th[i] = th[i-1] + (k1th + 2*k2th + 2*k3th + k4th)/6

    return r, th, dr, dth, t

# Solve the equations
r, th, dr, dth, t = rk4(r0, th0, dr0, dth0, mu, g, dt, t)

# Convert polar coordinates to Cartesian coordinates
x = r*np.sin(th)
y = -r*np.cos(th)

# Plot the trajectory in the xy plane
plt.figure(figsize=(7, 6))
plt.plot(x, y, 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()

# Define the cutting plane for the Poincaré section
cut = 0  # e.g., cutting angle at theta = 0 rad

# Identify the intersections of trajectories with the cutting plane
poincare_points = []
for i in range(1, len(th)):
    # Check if theta crosses the cutting angle in the positive direction
    if (th[i-1] < cut and th[i] >= cut and dth[i] > 0) or \
       (th[i-1] > cut and th[i] <= cut and dth[i] < 0):
        poincare_points.append((r[i], dr[i]))

# Convert Poincaré points to an array for plotting
poincare_points = np.array(poincare_points)

# Plot the Poincaré section
plt.figure(figsize=(8, 6))
plt.plot(poincare_points[:, 0], poincare_points[:, 1], 'bo', markersize=2)
plt.xlabel('$r \, (m)$', fontsize=12)
plt.ylabel('$\\dot{r} \, (rad/s)$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.tight_layout()
plt.show()