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()
 

aW1wb3J0IG51bXB5IGFzIG5wCmltcG9ydCBzY2lweS5vcHRpbWl6ZSBhcyBvcHQKaW1wb3J0IG1hdHBsb3RsaWIucHlwbG90IGFzIHBsdAppbXBvcnQgcGFuZGFzIGFzIHBkCgojIERlZmluZSB1cGRhdGVkIHBhcmFtZXRlcnMKbXUgPSAwLjM1OTExNzgwNDcgICAgICAjIE1hc3MgcmF0aW8gb2YgdGhlIHNlY29uZGFyeSBib2R5Cm11MyA9IDY2OTYuNDI0ODE0ICAgICAgIyBNYXNzIHJhdGlvIG9mIHRoZSB0ZXJ0aWFyeSBib2R5ClIzID0gNDMuNzYzOTg0ICAgICAgICAgIyBEaXN0YW5jZSBvZiB0aGUgdGVydGlhcnkgYm9keSAoY2Fub25pY2FsIHVuaXRzKQoKIyBEZWZpbmUgc2hpIGFzIGEgZnVuY3Rpb24gb2YgdGltZQpkZWYgc2hpX2Z1bmN0aW9uKHQpOgogICAgc2hpXzAgPSBucC5waSAvIDIKICAgIGRlbHRhX24gPSAxIC0gbjMgICMgRGlmZmVyZW5jZSBpbiBhbmd1bGFyIHNwZWVkCiAgICByZXR1cm4gc2hpXzAgKyBkZWx0YV9uICogdAoKIyBDb3JyZWN0ZWQgYW5ndWxhciBzcGVlZCBvZiBtMyBpbiBjYW5vbmljYWwgdW5pdHMgKGFzc3VtaW5nIHNoaSA9IM+ALzIgaW5pdGlhbGx5KQpzaGlfaW5pdGlhbCA9IG5wLnBpIC8gMgpuMyA9IG5wLnNxcnQoKDEgLSBtdSkgLyBSMyoqMyArIG11IC8gKFIzKioyICsgMSAtIDIgKiBSMyAqIG5wLmNvcyhzaGlfaW5pdGlhbCkpKiooMy8yKSkgICMgVXBkYXRlZCBmb3JtdWxhCgojIERlZmluZSBmdW5jdGlvbnMgZm9yIGRpc3RhbmNlcyAoa2VlcGluZyB4ID0gMCkKZGVmIHIxKHksIHopOgogICAgcmV0dXJuIG5wLnNxcnQobXUgKiogMiArIHkgKiogMiArIHogKiogMikKCmRlZiByMih5LCB6KToKICAgIHJldHVybiBucC5zcXJ0KCgxIC0gbXUpICoqIDIgKyB5ICoqIDIgKyB6ICoqIDIpCgpkZWYgcjMoeSwgeiwgc2hpKToKICAgIHJldHVybiBucC5zcXJ0KChSMyAqIG5wLmNvcyhzaGkpIC0gbXUpICoqIDIgKyB5ICoqIDIgKyB6ICoqIDIpCgojIERlZmluZSBmb3JjZSBiYWxhbmNlIGVxdWF0aW9ucyB3aXRoIGNvcnJlY3Rpb24gdGVybXMgKHggPSAwKQpkZWYgbGFncmFuZ2VfZXF1YXRpb25zKHZhcnMsIHNoaSk6CiAgICB5LCB6ID0gdmFycwoKICAgICMgQ29tcHV0ZSBkaXN0YW5jZXMKICAgIHIxX3ZhbCA9IHIxKHksIHopCiAgICByMl92YWwgPSByMih5LCB6KQogICAgcjNfdmFsID0gcjMoeSwgeiwgc2hpKQoKICAgICMgQ29ycmVjdGlvbiB0ZXJtIChvbmx5IGluIHkgZGlyZWN0aW9uLCBzaW5jZSB4PTApCiAgICBjb3JyZWN0aW9uX3kgPSBtdTMgKiAoKDEgLSBtdSkgKiBucC5zaW4oc2hpKSAvIFIzKioyICsgbXUgKiAoUjMgKiBucC5zaW4oc2hpKSkgLyAoUjMqKjIgKyAxIC0gMiAqIFIzICogbnAuY29zKHNoaSkpKiooMy8yKSkKCiAgICAjIEZvcmNlIGVxdWF0aW9ucyAoRnggaXMgcmVtb3ZlZCBiZWNhdXNlIHggPSAwKQogICAgRnkgPSAtKDEgLSBtdSkgKiB5IC8gcjFfdmFsICoqIDMgLSBtdSAqIHkgLyByMl92YWwgKiogMyAtIG11MyAqIHkgLyByM192YWwgKiogMyAtIGNvcnJlY3Rpb25feSAtIHkKICAgIEZ6ID0gLSgxIC0gbXUpICogeiAvIHIxX3ZhbCAqKiAzIC0gbXUgKiB6IC8gcjJfdmFsICoqIDMgLSBtdTMgKiB6IC8gcjNfdmFsICoqIDMKCiAgICByZXR1cm4gW0Z5LCBGel0KCiMgR2VuZXJhdGUgaW5pdGlhbCBndWVzc2VzIGluIHkteiBwbGFuZSAoc2V0dGluZyB4PTApCmluaXRpYWxfZ3Vlc3NlcyA9IFsKICAgICgwLjUsIDApLCAoLTAuNSwgMCksICAjIEV4cGxvcmluZyBwb2ludHMgYWxvbmcgeS1heGlzCiAgICAoMCwgMC41KSwgKDAsIC0wLjUpLCAoMC41LCAwLjUpLCAoMC41LCAtMC41KSwgKC0wLjUsIDAuNSksICgtMC41LCAtMC41KSwgICMgRXhwbG9yaW5nIGRpZmZlcmVudCB5LXogbG9jYXRpb25zCiAgICAoMC41LCAxKSwgKDAuNSwgLTEpLCAoLTAuNSwgMSksICgtMC41LCAtMSkKXQoKIyBDb21wdXRlIExhZ3JhbmdlIHBvaW50cyBmb3IgdCA9IDAgYW5kIHQgPSDPgC8yCnRpbWVfdmFsdWVzID0gWzAsIG5wLnBpIC8gMl0KcmVzdWx0cyA9IHt9Cgpmb3IgdCBpbiB0aW1lX3ZhbHVlczoKICAgIHNoaV90ID0gc2hpX2Z1bmN0aW9uKHQpCiAgICBudW1lcmljYWxfc29sdXRpb25zID0gW10KICAgIAogICAgZm9yIGd1ZXNzIGluIGluaXRpYWxfZ3Vlc3NlczoKICAgICAgICBzb2wgPSBvcHQuZnNvbHZlKGxhZ3JhbmdlX2VxdWF0aW9ucywgZ3Vlc3MsIGFyZ3M9KHNoaV90LCkpCiAgICAgICAgbnVtZXJpY2FsX3NvbHV0aW9ucy5hcHBlbmQoc29sKQoKICAgICMgQ29udmVydCB0byBEYXRhRnJhbWUgKFJlbW92aW5nIGR1cGxpY2F0ZXMpCiAgICBkZiA9IHBkLkRhdGFGcmFtZShudW1lcmljYWxfc29sdXRpb25zLCBjb2x1bW5zPVsneScsICd6J10pCiAgICBkZiA9IGRmLnJvdW5kKDYpLmRyb3BfZHVwbGljYXRlcygpLnJlc2V0X2luZGV4KGRyb3A9VHJ1ZSkgICMgUmVtb3ZlIGR1cGxpY2F0ZXMgYW5kIHJvdW5kCiAgICByZXN1bHRzW3RdID0gZGYKCiAgICAjIERpc3BsYXkgcmVzdWx0cwogICAgcHJpbnQoZiJcbk51bWVyaWNhbCBMYWdyYW5nZSBQb2ludHMgZm9yIHQgPSB7dH0gKHkteiBQbGFuZSwgeD0wKToiKQogICAgcHJpbnQoZGYpCgogICAgIyBQbG90IHRoZSByZXN1bHRzCiAgICBmaWcgPSBwbHQuZmlndXJlKGZpZ3NpemU9KDcsIDcpKQogICAgYXggPSBmaWcuYWRkX3N1YnBsb3QoMTExKQoKICAgIGF4LnNjYXR0ZXIoZGZbJ3knXSwgZGZbJ3onXSwgY29sb3I9J3JlZCcsIGxhYmVsPSdMYWdyYW5nZSBQb2ludHMnKQogICAgYXguc2NhdHRlcigwLCAwLCBjb2xvcj0nYmx1ZScsIG1hcmtlcj0nbycsIGxhYmVsPSdQcmltYXJ5IEJvZHknKQogICAgYXguc2NhdHRlcigwLCAxLCBjb2xvcj0nZ3JlZW4nLCBtYXJrZXI9J28nLCBsYWJlbD0nU2Vjb25kYXJ5IEJvZHknKQoKICAgIGF4LnNldF94bGFiZWwoIlkiKQogICAgYXguc2V0X3lsYWJlbCgiWiIpCiAgICBheC5sZWdlbmQoKQogICAgYXguc2V0X3RpdGxlKGYiTGFncmFuZ2UgUG9pbnRzIGluIHkteiBQbGFuZSAodCA9IHt0fSwgeD0wKSIpCiAgICBwbHQuc2hvdygpCg==