"""Problem 4: the Mars Orbiter Mission, reconstructed.

Left: the seven Earth-bound orbits, to scale (apse lines idealised as aligned).
Middle: the heliocentric transfer - the Lambert arc between the Earth on 30 Nov 2013 and Mars on
24 Sep 2014 (plan94 positions), against a Hohmann ellipse started on the same day.
Right: departure plus arrival excess speed over the 2013 window, with MOM's dates marked.
Run from the sfm-slides directory:   python3 codes/fig-im-mom.py
"""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, PURPLE
from _imcommon import (Re, AU, muS, DAY, MOM_EARTH, sma, planet, lambert, JD_TMI, JD_MOI)

fig = plt.figure(figsize=(17, 5.8))
aL = fig.add_subplot(1, 3, 1)
aM = fig.add_subplot(1, 3, 2)
aR = fig.add_subplot(1, 3, 3)
t = np.linspace(0, 2 * np.pi, 400)
cols = [BLUE, GREEN, AMBER, PURPLE, MAROON, "#888888", INK]
for o, c in zip(MOM_EARTH, cols):
    a = sma(*o)
    e = (o[1] - o[0]) / (2 * a)
    f = np.linspace(-np.pi, np.pi, 600)
    r = a * (1 - e * e) / (1 + e * np.cos(f))
    aL.plot(r * np.cos(f) / 1e3, r * np.sin(f) / 1e3, color=c, lw=1.6, label=f"{o[0]:.0f} x {o[1]:,.0f} km")
aL.fill(Re * np.cos(t) / 1e3, Re * np.sin(t) / 1e3, color=BLUE, alpha=0.5)
aL.set_aspect("equal")
aL.set_xlim(-205, 25)
aL.set_ylim(-120, 70)
aL.set_xlabel("[1000 km]")
aL.legend(fontsize=8.5, frameon=False, loc="lower center", title="launch orbit, then after each burn (6 to 16 Nov 2013)",
          title_fontsize=8.5, ncol=2)
aL.set_title("Earth-bound: six perigee burns (one fell short)", fontsize=11.5)

rE, vE = planet(JD_TMI, 3)
rMa, vMa = planet(JD_MOI, 4)
v1, v2 = lambert(rE, rMa, (JD_MOI - JD_TMI) * DAY)
# propagate the Lambert arc (two-body, Sun) for drawing
from scipy.integrate import solve_ivp
arc = solve_ivp(lambda tt, s: [*s[3:], *(-muS * s[:3] / np.linalg.norm(s[:3]) ** 3)],
                [0, (JD_MOI - JD_TMI) * DAY], [*rE, *v1], rtol=1e-10, atol=1e-3, dense_output=True)
ta = np.linspace(0, (JD_MOI - JD_TMI) * DAY, 400)
X = arc.sol(ta)
for k, (name, c) in enumerate((("Earth", BLUE), ("Mars", MAROON))):
    jj = JD_TMI + np.linspace(0, 690 if k else 366, 400)
    P = np.array([planet(j, 3 + k)[0] for j in jj])
    aM.plot(P[:, 0] / AU, P[:, 1] / AU, color=c, lw=0.9, ls=(0, (4, 3)))
aM.plot(X[0] / AU, X[1] / AU, color=GREEN, lw=2.2, label=f"Lambert arc, {JD_MOI - JD_TMI:.0f} d")
r1 = np.linalg.norm(rE)
ah = (r1 + 1.52368 * AU) / 2
eh = 1 - r1 / ah
th0 = np.arctan2(rE[1], rE[0])
f = np.linspace(0, np.pi, 300)
rh = ah * (1 - eh ** 2) / (1 + eh * np.cos(f))
aM.plot(rh * np.cos(f + th0) / AU, rh * np.sin(f + th0) / AU, color=INK, lw=1.2, ls=(0, (2, 2)),
        label="Hohmann ellipse from the same day, 259 d")
aM.scatter([rE[0] / AU], [rE[1] / AU], color=BLUE, zorder=4)
aM.text(rE[0] / AU + 0.05, rE[1] / AU, "Earth at TMI\n30 Nov 2013", fontsize=9, color=BLUE)
aM.scatter([rMa[0] / AU], [rMa[1] / AU], color=MAROON, zorder=4)
aM.text(rMa[0] / AU + 0.05, rMa[1] / AU - 0.1, "Mars at MOI\n24 Sep 2014", fontsize=9, color=MAROON)
aM.scatter([0], [0], color=AMBER, s=120, zorder=4)
aM.set_aspect("equal")
aM.set_xlabel("ecliptic $x$ [AU]")
aM.set_ylabel("ecliptic $y$ [AU]")
aM.legend(fontsize=8.5, frameon=False, loc="lower left")
aM.set_title("Heliocentric transfer, reconstructed", fontsize=11.5)

dep = np.arange(-40, 41, 2.0)
tof = np.arange(190, 381, 3.0)
Z = np.full((len(tof), len(dep)), np.nan)
for i, dd in enumerate(dep):
    rE_, vE_ = planet(JD_TMI + dd, 3)
    for j, tt in enumerate(tof):
        rM_, vM_ = planet(JD_TMI + dd + tt, 4)
        a_, b_ = lambert(rE_, rM_, tt * DAY)
        Z[j, i] = np.linalg.norm(a_ - vE_) + np.linalg.norm(b_ - vM_)
cs = aR.contourf(dep, tof, Z, levels=np.arange(5.8, 9.01, 0.2), cmap="viridis_r", extend="max")
fig.colorbar(cs, ax=aR, label=r"$v_{\infty,\rm dep}+v_{\infty,\rm arr}$ [km/s]")
aR.scatter([0], [JD_MOI - JD_TMI], marker="*", s=220, color=MAROON, edgecolor="white", zorder=4, label="MOM")
aR.set_xlabel("departure, days after 30 Nov 2013")
aR.set_ylabel("time of flight [d]")
aR.legend(loc="upper right", framealpha=0.9)
aR.set_title("The 2013 window: total excess speed", fontsize=11.5)
fig.tight_layout()
fig.savefig("img/fig-im-mom.png", dpi=150, bbox_inches="tight")
print("wrote img/fig-im-mom.png")