"""Problem 5: Aditya-L1, reconstructed.

Left: the launch orbit and the four Earth-bound orbits, to scale (apse lines idealised as aligned).
Right: a CR3BP transfer from a 256 km perigee to the halo orbit, in the Sun-Earth rotating frame
(projection on the ecliptic), chosen from codes/data/aditya-hoi-*.npy (scan and refinement) as the cheapest
insertion among transfers of 105-120 days (ISRO's cruise took about 110 days).
Run from the sfm-slides directory:   python3 codes/fig-im-aditya.py
"""
import glob
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, PURPLE
from _imcommon import (Re, AU, VU, TU, DAY, MU3, AL1_EARTH, sma, cr3bp, halo, lagrange_L1, propagate)

fig, (aL, aR) = plt.subplots(1, 2, figsize=(14.5, 6.2), gridspec_kw=dict(width_ratios=[1, 1.15]))
t = np.linspace(0, 2 * np.pi, 400)
dates = ["2 Sep (PSLV)", "3 Sep", "5 Sep", "10 Sep", "15 Sep"]
for o, c, d in zip(AL1_EARTH, [BLUE, GREEN, AMBER, PURPLE, MAROON], dates):
    a = sma(*o)
    e = (o[1] - o[0]) / (2 * a)
    f = np.linspace(-np.pi, np.pi, 600)
    r = a * (1 - e * e) / (1 + e * np.cos(f))
    aL.plot(r * np.cos(f) / 1e3, r * np.sin(f) / 1e3, color=c, lw=1.8, label=f"{d}: {o[0]:.0f} x {o[1]:,.0f} km")
aL.fill(Re * np.cos(t) / 1e3, Re * np.sin(t) / 1e3, color=BLUE, alpha=0.5)
aL.set_aspect("equal")
aL.set_xlim(-135, 15)
aL.set_ylim(-75, 50)
aL.set_xlabel("[1000 km]")
aL.legend(fontsize=9, frameon=False, loc="lower center", ncol=2, title="orbit after each burn (2023)",
          title_fontsize=9)
aL.set_title("Earth-bound phase: four perigee burns in 13 days", fontsize=12)

xL1 = lagrange_L1()
xE = 1 - MU3
s0, T = halo(120000.0)
ts, H = propagate(s0, T, 1500)
crossings = {0: s0, 1: solve_ivp(cr3bp, [0, T / 2], s0, rtol=1e-12, atol=1e-13).y[:, -1]}
scan = np.vstack([np.load(f)[:, :5] for f in sorted(glob.glob("codes/data/aditya-hoi-*.npy"))])
sel = scan[(scan[:, 1] > 105) & (scan[:, 1] < 120)]
best = sel[np.argmin(sel[:, 2])]
k, tof, dv, th, ph = int(best[0]), best[1], best[2], best[3], best[4]
d = np.array([np.sin(th) * np.cos(ph), np.sin(th) * np.sin(ph), np.cos(th)])
y0 = crossings[k].copy()
y0[3:] += dv / 1000 / VU * d
tr = solve_ivp(cr3bp, [0, -tof * DAY / TU], y0, rtol=1e-10, atol=1e-12, dense_output=True)
tt = np.linspace(0, -tof * DAY / TU, 2000)
X = tr.sol(tt)
kmx = lambda x: (x - xE) * AU / 1e6          # x relative to the Earth, in million km
aR.plot(kmx(H[0]), H[1] * AU / 1e6, color=MAROON, lw=2, label="halo orbit (CR3BP), period 177.9 d")
aR.plot(kmx(X[0]), X[1] * AU / 1e6, color=GREEN, lw=1.6, label=f"transfer from a 256 km perigee, {tof:.0f} d")
for day in range(10, int(tof), 10):
    p = tr.sol(-(tof - day) * DAY / TU) if False else tr.sol(-tof * DAY / TU + day * DAY / TU)
    aR.scatter(kmx(p[0]), p[1] * AU / 1e6, s=8, color=GREEN)
aR.scatter([0], [0], s=80, color=BLUE, zorder=4)
aR.text(0.03, 0.04, "Earth", color=BLUE, fontsize=10)
aR.scatter([kmx(xL1)], [0], marker="+", s=120, color=INK, zorder=4)
aR.text(kmx(xL1) + 0.03, 0.04, "$L_1$", fontsize=12)
aR.scatter([kmx(y0[0])], [y0[1] * AU / 1e6], marker="*", s=160, color=AMBER, zorder=5,
           label=f"insertion at the xz-plane crossing, $\\Delta v={dv:.0f}$ m/s")
aR.annotate("to the Sun", xy=(-1.95, -0.75), xytext=(-1.55, -0.75), arrowprops=dict(arrowstyle="->", color=AMBER),
            color=AMBER, va="center")
aR.set_aspect("equal")
aR.set_xlabel("$x$ from the Earth (towards the Sun is negative) [million km]")
aR.set_ylabel("$y$ [million km]")
aR.legend(fontsize=8.5, frameon=False, loc="upper left")
aR.set_title("Cruise and halo, Sun-Earth rotating frame (dots every 10 days)", fontsize=12)
fig.tight_layout()
fig.savefig("img/fig-im-aditya.png", dpi=150, bbox_inches="tight")
print(f"wrote img/fig-im-aditya.png  (transfer: crossing {'AB'[k]}, {tof:.1f} d, HOI {dv:.1f} m/s)")