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

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

# Parameters
mu = 16      # reduced mass 

# Initial conditions 
#(a) (1, 1.4, 0, 0, 1.527) 
#(b) (0.25, 0.35, 0, 2.83, 3) 
#(c) (1, 0, np.pi, 1, 4.35) 
#(d) (0.5, 0, np.pi, 1, 1.95)

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

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

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

def rk4(r0, th0, dr0, dth0, mu, 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)
        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)
        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)
        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)
        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, 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)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
# 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, th, 'b-', linewidth=2)
plt.xlabel('$t \, (s)$', fontsize=14)
plt.ylabel('$\\theta \, (rad)$', fontsize=14)
plt.grid(True)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)

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

# Angular phase space plot
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.plot(th, dth, 'b-', linewidth=2)
plt.xlabel('$\\theta \, (rad)$', fontsize=14)
plt.ylabel('$\\dot{\\theta} \, (rad/s)$', fontsize=14)
plt.grid(True)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)

# Radial phase space plot
plt.subplot(1, 2, 2)
plt.plot(r, dr, 'r-', linewidth=2)
plt.xlabel('$r \, (rad)$', fontsize=14)
plt.ylabel('$\\dot{r} \, (rad/s)$', fontsize=14)
plt.grid(True)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.tight_layout()
plt.show()

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