"""Problem 3: Chandrayaan-3 at the Moon, and its propulsion module's return to Earth orbit.

Left: the lunar orbits from insertion to 153 x 163 km, and the module's departure (150 km circle,
then 150 x 5112 km, then the escape hyperbola), to scale.
Right: the Earth orbit the module was moved to, with the Moon's orbit and the GEO belt.
Reconstructed from ISRO's published figures; apse lines idealised.
Run from the sfm-slides directory:   python3 codes/fig-im-chandrayaan3.py
"""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, PURPLE
from _imcommon import Re, RM, rM, muM, mu, CH3_MOON, sma, vis, DAY


def ellipse(p_alt, a_alt, R, w=0.0, n=900):
    a = sma(p_alt, a_alt, R)
    e = (a_alt - p_alt) / (2 * a)
    f = np.linspace(-np.pi, np.pi, n)
    r = a * (1 - e * e) / (1 + e * np.cos(f))
    return r * np.cos(f + w), r * np.sin(f + w)


t = np.linspace(0, 2 * np.pi, 400)
fig, (aL, aR) = plt.subplots(1, 2, figsize=(14.5, 6.6), gridspec_kw=dict(width_ratios=[1, 1.05]))
dates = ["5 Aug", "6 Aug", "9 Aug", "14 Aug", "16 Aug"]
for o, c, d in zip(CH3_MOON, [MAROON, AMBER, GREEN, PURPLE, BLUE], dates):
    x, y = ellipse(*o, RM, w=-np.pi / 2)
    aL.plot(x / 1e3, y / 1e3, color=c, lw=1.6, label=f"{d}: {o[0]} x {o[1]:,} km")
# the module's departure: 150 x 5112 km (9 Oct), then the escape hyperbola (13 Oct)
x, y = ellipse(150, 5112, RM, w=np.pi / 2)
aL.plot(x / 1e3, y / 1e3, color=INK, lw=1.6, ls=(0, (5, 3)), label="9 Oct: 150 x 5,112 km (module)")
vinf = 0.20
rp = RM + 150
e = 1 + rp * vinf ** 2 / muM
p = rp * (1 + e)
f = np.linspace(-0.2, np.arccos(-1 / e) - 0.05, 400)
r = p / (1 + e * np.cos(f))
aL.plot(r * np.cos(f + np.pi / 2) / 1e3, r * np.sin(f + np.pi / 2) / 1e3, color=INK, lw=2,
        label=r"13 Oct: escape ($v_\infty\approx0.2$ km/s, sketch)")
aL.fill(RM * np.cos(t) / 1e3, RM * np.sin(t) / 1e3, color="#9a9388", alpha=0.6)
aL.set_aspect("equal")
aL.set_xlim(-14, 14)
aL.set_ylim(-21, 9)
aL.set_xlabel("[1000 km]")
aL.legend(fontsize=8.5, frameon=False, loc="lower left", title="Moon-centred orbits (perilune x apolune altitude)",
          title_fontsize=8.5)
aL.set_title("At the Moon: capture, lowering, and the module's departure", fontsize=12)

aR.plot(rM * np.cos(t) / 1e3, rM * np.sin(t) / 1e3, color=INK, lw=0.9, ls=(0, (5, 4)), label="Moon's orbit")
aR.plot(42164 * np.cos(t) / 1e3, 42164 * np.sin(t) / 1e3, color=AMBER, lw=1.4, label="GEO belt")
for (pa, aa), c, lab in (((180000, 380000), GREEN, "target, 1.8 lakh x 3.8 lakh km"),
                         ((154000, None), MAROON, "achieved: perigee 1.54 lakh km, period about 13 d")):
    if aa is None:                      # apogee from the period
        a = (mu * (13 * DAY) ** 2 / (4 * np.pi ** 2)) ** (1 / 3)
        aa = 2 * a - 2 * Re - pa
    x, y = ellipse(pa, aa, Re, w=np.radians(30))
    aR.plot(x / 1e3, y / 1e3, color=c, lw=2, label=f"{lab} ({pa / 1e3:.0f} x {aa / 1e3:.0f} thousand km)")
aR.fill(Re * 3 * np.cos(t) / 1e3, Re * 3 * np.sin(t) / 1e3, color=BLUE, alpha=0.5)
aR.set_aspect("equal")
aR.set_xlabel("[1000 km]")
aR.set_xlim(-480, 480)
aR.set_ylim(-480, 480)
aR.legend(fontsize=8.5, frameon=False, loc="lower left")
aR.set_title("Back at the Earth: the module's new orbit (Earth drawn 3x)", fontsize=12)
fig.tight_layout()
fig.savefig("img/fig-im-chandrayaan3.png", dpi=150, bbox_inches="tight")
print("wrote img/fig-im-chandrayaan3.png")