"""Problem 1: Chandrayaan-1's orbits, reconstructed from ISRO's published perigee and apogee figures.

Left: the five Earth-bound orbits and the lunar transfer, to scale (apse lines idealised as aligned).
Right: the lunar orbits from insertion to the 100 km circle, to scale.
Run from the sfm-slides directory:   python3 codes/fig-im-chandrayaan1.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, CH1_EARTH, CH1_MOON, CH1_DATES, sma


def ellipse(p_alt, a_alt, R, w=0.0, n=800):
    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)


cols = [BLUE, GREEN, AMBER, PURPLE, MAROON, INK]
fig, (aL, aR) = plt.subplots(1, 2, figsize=(14.5, 6.4), gridspec_kw=dict(width_ratios=[1.45, 1]))
t = np.linspace(0, 2 * np.pi, 400)
aL.plot(rM * np.cos(t) / 1e3, rM * np.sin(t) / 1e3, color=INK, lw=0.8, ls=(0, (5, 4)), alpha=0.6)
aL.text(-rM / 1e3 + 6, 100, "Moon's\norbit", fontsize=10, color=INK)
for k, (o, c) in enumerate(zip(CH1_EARTH, cols)):
    x, y = ellipse(*o, Re, w=np.pi)                       # perigee to the right, apogee to the left
    x, y = -x, y
    lab = f"{CH1_DATES[k]}: {o[0]:.0f} x {o[1]:,.0f} km" + (" (lunar transfer)" if k == 5 else "")
    aL.plot(x / 1e3, y / 1e3, color=c, lw=1.8 if k < 5 else 2.2, ls="-" if k < 5 else (0, (6, 3)), label=lab)
aL.fill(Re * np.cos(t) / 1e3, Re * np.sin(t) / 1e3, color=BLUE, alpha=0.5)
aL.scatter([-rM / 1e3], [0], s=60, color=INK, zorder=4)
aL.text(-rM / 1e3, -22, "Moon at\narrival", ha="center", va="top", fontsize=10)
aL.set_aspect("equal")
aL.set_xlabel("$x$ [1000 km]")
aL.set_ylabel("$y$ [1000 km]")
aL.legend(loc="upper center", bbox_to_anchor=(0.5, -0.14), ncol=2, fontsize=9, frameon=False,
          title="orbit after each burn (perigee x apogee altitude)", title_fontsize=9)
aL.set_title("Earth-bound phase: five perigee burns, 22 Oct to 4 Nov 2008", fontsize=12)
aL.set_xlim(-410, 60)
aL.set_ylim(-120, 120)

for k, (o, c) in enumerate(zip(CH1_MOON, [MAROON, AMBER, GREEN, PURPLE, BLUE])):
    x, y = ellipse(*o, RM, w=-np.pi / 2)                  # perilune at the bottom
    aR.plot(x / 1e3, y / 1e3, color=c, lw=1.8, label=f"{o[0]} x {o[1]:,} km")
aR.fill(RM * np.cos(t) / 1e3, RM * np.sin(t) / 1e3, color="#9a9388", alpha=0.6)
aR.set_aspect("equal")
aR.set_xlabel("[1000 km]")
aR.legend(loc="upper right", fontsize=9, frameon=False, title="lunar orbits, 8 to 12 Nov", title_fontsize=9)
aR.set_title("Lunar phase: capture, then four lowering burns", fontsize=12)
aR.set_xlim(-6, 6)
aR.set_ylim(-3.5, 10.8)
fig.tight_layout()
fig.savefig("img/fig-im-chandrayaan1.png", dpi=150, bbox_inches="tight")
print("wrote img/fig-im-chandrayaan1.png")