import numpy as np
import scipy.optimize as opt
import matplotlib.pyplot as plt
import pandas as pd

# Define updated parameters
mu = 0.3591178047      # Mass ratio of the secondary body
mu3 = 6696.424814      # Mass ratio of the tertiary body
R3 = 43.763984         # Distance of the tertiary body (canonical units)

# Define shi as a function of time
def shi_function(t):
    shi_0 = np.pi / 2
    delta_n = 1 - n3  # Difference in angular speed
    return shi_0 + delta_n * t

# Corrected angular speed of m3 in canonical units (assuming shi = π/2 initially)
shi_initial = np.pi / 2
n3 = np.sqrt((1 - mu) / R3**3 + mu / (R3**2 + 1 - 2 * R3 * np.cos(shi_initial))**(3/2))  # Updated formula

# Define functions for distances (keeping x = 0)
def r1(y, z):
    return np.sqrt(mu ** 2 + y ** 2 + z ** 2)

def r2(y, z):
    return np.sqrt((1 - mu) ** 2 + y ** 2 + z ** 2)

def r3(y, z, shi):
    return np.sqrt((R3 * np.cos(shi) - mu) ** 2 + y ** 2 + z ** 2)

# Define force balance equations with correction terms (x = 0)
def lagrange_equations(vars, shi):
    y, z = vars

    # Compute distances
    r1_val = r1(y, z)
    r2_val = r2(y, z)
    r3_val = r3(y, z, shi)

    # Correction term (only in y direction, since x=0)
    correction_y = mu3 * ((1 - mu) * np.sin(shi) / R3**2 + mu * (R3 * np.sin(shi)) / (R3**2 + 1 - 2 * R3 * np.cos(shi))**(3/2))

    # Force equations (Fx is removed because x = 0)
    Fy = -(1 - mu) * y / r1_val ** 3 - mu * y / r2_val ** 3 - mu3 * y / r3_val ** 3 - correction_y - y
    Fz = -(1 - mu) * z / r1_val ** 3 - mu * z / r2_val ** 3 - mu3 * z / r3_val ** 3

    return [Fy, Fz]

# Generate initial guesses in y-z plane (setting x=0)
initial_guesses = [
    (0.5, 0), (-0.5, 0),  # Exploring points along y-axis
    (0, 0.5), (0, -0.5), (0.5, 0.5), (0.5, -0.5), (-0.5, 0.5), (-0.5, -0.5),  # Exploring different y-z locations
    (0.5, 1), (0.5, -1), (-0.5, 1), (-0.5, -1)
]

# Compute Lagrange points for t = 0 and t = π/2
time_values = [0, np.pi / 2]
results = {}

for t in time_values:
    shi_t = shi_function(t)
    numerical_solutions = []
    
    for guess in initial_guesses:
        sol = opt.fsolve(lagrange_equations, guess, args=(shi_t,))
        numerical_solutions.append(sol)

    # Convert to DataFrame (Removing duplicates)
    df = pd.DataFrame(numerical_solutions, columns=['y', 'z'])
    df = df.round(6).drop_duplicates().reset_index(drop=True)  # Remove duplicates and round
    results[t] = df

    # Display results
    print(f"\nNumerical Lagrange Points for t = {t} (y-z Plane, x=0):")
    print(df)

    # Plot the results
    fig = plt.figure(figsize=(7, 7))
    ax = fig.add_subplot(111)

    ax.scatter(df['y'], df['z'], color='red', label='Lagrange Points')
    ax.scatter(0, 0, color='blue', marker='o', label='Primary Body')
    ax.scatter(0, 1, color='green', marker='o', label='Secondary Body')

    ax.set_xlabel("Y")
    ax.set_ylabel("Z")
    ax.legend()
    ax.set_title(f"Lagrange Points in y-z Plane (t = {t}, x=0)")
    plt.show()