"""Q5: Chandrayaan-3 -- the orbit chain, and how much a 1 m/s burn error moves the apogee.

Left: the three orbits (to scale) with the Moon's orbit.  Right: apogee error against
burn-speed error for burns at different semi-major axes: the gain is 4 a^2 v / mu.

Run from the sfm-slides directory:   python3 codes/fig-mq-ch3-raising.py
"""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _mqcommon import mu, Re, vis, conic_xy

rp = Re + 200.0
A, B, C = Re + 36500, Re + 127603, Re + 369328
fig, (axL, axR) = plt.subplots(1, 2, figsize=(14.4, 5.8), gridspec_kw=dict(width_ratios=[1.25, 1]))

t = np.linspace(0, 2*np.pi, 600)
axL.plot(384400*np.cos(t), 384400*np.sin(t), color=INK, lw=1.0, ls=(0, (5, 4)), alpha=0.7)
axL.text(0, -396000, "Moon's orbit, 384 400 km", ha="center", fontsize=10, color=INK, alpha=0.8, va="top")
for ra, col, lab in ((A, BLUE, "after launch:  170 $\\times$ 36 500 km"),
                     (B, GREEN, "after 5 burns:  236 $\\times$ 127 603 km"),
                     (C, MAROON, "after TLI:  288 $\\times$ 369 328 km")):
    a = (rp + ra)/2; e = (ra - rp)/(ra + rp)
    x, y = conic_xy(a, e)
    axL.plot(x, y, color=col, lw=2.2, label=lab)
axL.fill(Re*np.cos(t)*6, Re*np.sin(t)*6, color=BLUE, alpha=0.18)         # Earth drawn 6x for visibility
axL.plot(rp, 0, "o", color=INK, ms=7, zorder=6)
axL.annotate("every burn is at perigee", (rp, 0), (60000, 100000), fontsize=10, color=INK,
             arrowprops=dict(arrowstyle="-", color=INK, lw=0.8))
axL.set_aspect("equal"); axL.axis("off")
axL.set_xlim(-420000, 120000); axL.set_ylim(-420000, 420000)
axL.legend(frameon=False, fontsize=10.5, loc="lower right", bbox_to_anchor=(1.02, 0.02))
axL.set_title("Chandrayaan-3: the orbit chain (Earth drawn 6x too large)", fontsize=11.5)

# ---- B : the gain 4 a^2 v / mu ----------------------------------------------------
dvm = np.linspace(0, 3, 50)
cases = (("LEO $\\to$ GEO transfer, departure burn   ($a=24\\,421$ km)", 24421.07, 10.15149, BLUE),
         ("EBN: 36 500 $\\to$ 127 603 km   ($a\\approx 70\\,000$ km)", (rp + B)/2, vis(rp, (rp + B)/2), GREEN),
         ("TLI: 127 603 $\\to$ 369 328 km   ($a\\approx 190\\,000$ km)", (rp + C)/2, vis(rp, (rp + C)/2), MAROON))
for lab, a, v, col in cases:
    g = 4*a**2*v/mu/1000                       # km of apogee radius per m/s
    axR.plot(dvm, g*dvm, color=col, lw=2.5, label=f"{lab}:  {g:,.0f} km per m/s")
axR.axhspan(0, 1737, color=AMBER, alpha=0.18)
axR.text(0.05, 1250, "inside one lunar radius (1737 km)", fontsize=10, color=AMBER, va="center")
axR.set_xlabel(r"error in the burn, $\delta(\Delta v)$  [m/s]")
axR.set_ylabel(r"error in the apogee radius, $\delta r_a$  [km]")
axR.set_xlim(0, 3); axR.set_ylim(0, 13000)
axR.legend(frameon=False, fontsize=9.8, loc="upper left")
axR.set_title(r"Sensitivity $\delta r_a/\delta(\Delta v)=4a^{2}v/\mu$ grows with the orbit", fontsize=11.5)
despine(axR)

plt.tight_layout()
fig.savefig("img/fig-mq-ch3-raising.png", dpi=160, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-mq-ch3-raising.png")
