"""Problem 2: the Earth-bound phases of Chandrayaan-1, -2 and -3 side by side.

Left: apogee altitude against days from launch (each step is one burn; the last is the trans-lunar
injection).  Right: the impulsive dv of each burn, from the published orbits.
Run from the sfm-slides directory:   python3 codes/fig-im-chandrayaans.py
"""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import INK, MAROON, BLUE, GREEN, despine
from _imcommon import CH1_EARTH, CH2_EARTH, CH3_EARTH, chain

DAYS = {   # days from launch of each burn (the last one is TLI), from the dates ISRO announced
    "Chandrayaan-1 (PSLV-XL, 2008)": (CH1_EARTH, [1, 3, 4, 7, 13], BLUE),
    "Chandrayaan-2 (GSLV Mk III, 2019)": (CH2_EARTH, [2, 3, 7, 11, 15, 23], MAROON),
    "Chandrayaan-3 (LVM3, 2023)": (CH3_EARTH, [1, 3, 4, 6, 11, 18], GREEN),
}
fig, (aL, aR) = plt.subplots(1, 2, figsize=(14, 5.2))
for k, (name, (orb, days, c)) in enumerate(DAYS.items()):
    apo = [o[1] for o in orb]
    t = [0] + days
    aL.step(t + [t[-1] + 2], apo + [apo[-1]], where="post", color=c, lw=2, label=name)
    aL.scatter(t, apo, color=c, s=18, zorder=3)
    dv = [d * 1000 for d, _ in chain(orb) if not np.isnan(d)]       # trim burns left out
    x = np.arange(len(dv)) + (k - 1) * 0.27
    aR.bar(x, dv, width=0.26, color=c, label=f"{name.split(' (')[0]}: total {sum(dv):.0f} m/s")
aL.set_yscale("log")
aL.set_yticks([2e4, 5e4, 1e5, 2e5, 4e5])
aL.set_yticklabels(["20 000", "50 000", "100 000", "200 000", "400 000"])
aL.minorticks_off()
aL.axhline(35786, color=INK, lw=0.8, ls=":")
aL.text(19, 39000, "GEO altitude", fontsize=9)
aL.set_xlabel("days from launch")
aL.set_ylabel("apogee altitude [km]")
aL.set_title("Apogee raised in steps; the last step is the trans-lunar injection", fontsize=12)
aL.legend(frameon=False, fontsize=9, loc="lower right")
aR.set_xticks(range(5))
aR.set_xticklabels([f"raise {i + 1}" for i in range(4)] + ["raise 5 / TLI"])
aR.set_ylabel(r"impulsive $\Delta v$ [m/s]")
aR.set_title("Apogee-raising burns and TLI (trim burns omitted)", fontsize=12)
aR.legend(frameon=False, fontsize=9)
aR.axhline(0, color=INK, lw=0.6)
for ax in (aL, aR):
    despine(ax)
fig.tight_layout()
fig.savefig("img/fig-im-chandrayaans.png", dpi=150, bbox_inches="tight")
print("wrote img/fig-im-chandrayaans.png")