"""Problem 6: Aditya-L1's halo orbit in the CR3BP, and the cost of inserting into it.

Left: the halo orbit through z = +120 000 km (differential correction), in the rotating frame
centred on L1, with its semi-axes.  Right: every transfer found by codes/aditya-hoi-scan.py
(256 km perigee to an xz-plane crossing of the halo): insertion dv against time of flight, with
the lower envelope; ISRO's cruise of about 110 days is marked.
Run from the sfm-slides directory:   python3 codes/fig-im-aditya-cr3bp.py
"""
import glob
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, despine
from _imcommon import AU, MU3, halo, lagrange_L1, propagate

xL1 = lagrange_L1()
s0, T = halo(120000.0)
ts, H = propagate(s0, T, 2000)
X, Y, Z = (H[0] - xL1) * AU / 1e3, H[1] * AU / 1e3, H[2] * AU / 1e3       # thousand km from L1

fig = plt.figure(figsize=(15, 6.2))
ax = fig.add_subplot(1, 2, 1, projection="3d")
ax.plot(X, Y, Z, color=MAROON, lw=2)
ax.plot(X, Y, 0 * Z - 150, color=INK, lw=0.7, alpha=0.4)
ax.scatter([0], [0], [0], marker="+", s=80, color=INK)
ax.text(0, 0, 15, "$L_1$")
ax.quiver(0, 0, -150, 300, 0, 0, color=BLUE, lw=1.2, arrow_length_ratio=0.12)
ax.text(320, 0, -150, "to the Earth", color=BLUE, fontsize=9)
ax.set_xlabel("$x$ [1000 km]")
ax.set_ylabel("$y$ [1000 km]")
ax.set_zlabel("$z$ [1000 km]")
Ax, Ay, Az = (X.max() - X.min()) / 2, Y.max(), Z.max()
ax.set_title(f"Halo orbit: $A_x={Ax:.0f}$, $A_y={Ay:.0f}$, $A_z={Az:.0f}$ thousand km;"
             f" period {T * 365.25636 / (2 * np.pi):.2f} d", fontsize=11)
ax.view_init(elev=22, azim=-60)
ax.set_box_aspect((1, 2.6, 0.75))

aR = fig.add_subplot(1, 2, 2)
scan = np.vstack([np.load(f)[:, :5] for f in sorted(glob.glob("codes/data/aditya-hoi-*.npy"))])
for k, (c, lab) in enumerate(((GREEN, "arrive at crossing A ($z=+120\\,000$ km)"), (AMBER, "arrive at crossing B"))):
    s = scan[scan[:, 0] == k]
    aR.scatter(s[:, 1], s[:, 2], s=14, color=c, alpha=0.7, label=lab)
edges = np.arange(np.floor(scan[:, 1].min() / 10) * 10, scan[:, 1].max() + 10, 10)
env = []
for a, b in zip(edges[:-1], edges[1:]):
    s = scan[(scan[:, 1] >= a) & (scan[:, 1] < b)]
    if len(s) >= 5:                              # only bins that were searched well
        env.append((0.5 * (a + b), s[:, 2].min()))
env = np.array(env)
aR.plot(env[:, 0], env[:, 1], color=INK, lw=1.5, marker="o", ms=4,
        label="cheapest found (10-day bins with at least 5 transfers)")
aR.axvspan(105, 115, color=MAROON, alpha=0.12)
aR.text(110, 15, "Aditya-L1:\nabout 110 d", ha="center", color=MAROON, fontsize=9)
aR.annotate("stable manifold:\nno insertion burn,\nabout 330 d", xy=(268, 4), xytext=(225, 40), fontsize=9,
            arrowprops=dict(arrowstyle="->", color=INK))
aR.set_xlim(20, 275)
aR.set_ylim(0, 410)
aR.set_xlabel("time of flight from perigee to insertion [d]")
aR.set_ylabel(r"halo-insertion $\Delta v$ [m/s]")
aR.set_title("The price of a short cruise", fontsize=12)
aR.legend(frameon=False, fontsize=9, loc="upper right")
despine(aR)
fig.tight_layout()
fig.savefig("img/fig-im-aditya-cr3bp.png", dpi=150, bbox_inches="tight")
print("wrote img/fig-im-aditya-cr3bp.png")
for c, d in env:
    print(f"   TOF {c:.0f} d: cheapest HOI {d:.1f} m/s")