"""Burning at periapsis to escape: what it saves, and what it costs in time.

Periapsis is the cheapest place to escape (codes/fig-man-escape-point.py), but
it is also the *lowest*, so one might expect to pay for the saving with a
longer climb out to the sphere of influence.  This figure asks by how much, and
the answer depends entirely on where the clock is started.

*From the burn*, periapsis is never slower.  After a prograde apsidal burn the
escape conic is a parabola whose periapsis q is the burn radius, and
integrating v_r = sqrt(2 mu / r) sqrt(1 - q/r) from q out to R gives

    t(q -> R) = (2/3) sqrt((R - q) / (2 mu)) (R + 2q),

        dt/dq  proportional to  (R - 2q),

so raising the burn point *lengthens* the trip for every q < R/2.  With
R = R_SOI = 9.25e5 km that threshold is 4.62e5 km, outside every apoapsis in
this family, so the deep burn reaches the SOI sooner as well as cheaper.

*From apoapsis*, the spacecraft must first coast down, and that half-period is
the real penalty.  It wins the argument at both ends and loses it in the
middle: the periapsis route is slower for e < 0.216 (where it saves almost no
dv anyway), faster for 0.216 < e < 0.939, and slower again above that, where
the orbit is so large that half a period outweighs the faster climb.

The family holds the periapsis fixed at LEO and varies e, i.e. "same parking
perigee, progressively more eccentric".

Run from the sfm-slides directory:

    python3 codes/fig-man-escape-penalty.py
"""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, despine

MU = 398600.4418          # km^3/s^2
RP = 6678.137             # km, LEO periapsis held fixed across the family

# Earth's sphere of influence, r_SOI = a_E (m_E/m_S)^(2/5)
AU_KM = 1.495978707e8
EARTH_SUN_MASS_RATIO = 3.0034896e-6
RSOI = AU_KM*EARTH_SUN_MASS_RATIO**0.4


def dv_escape(r, a, mu=MU):
    """Prograde impulse from an apsis at radius r on the ellipse of size a."""
    return np.sqrt(2.0*mu/r) - np.sqrt(mu*(2.0/r - 1.0/a))


def t_parabolic(q, R, mu=MU):
    """Time from the periapsis q of a parabola out to radius R (closed form)."""
    return (2.0/3.0)*np.sqrt((R - q)/(2.0*mu))*(R + 2.0*q)


e = np.linspace(0.0, 0.95, 2000)
a = RP/(1.0 - e)
ra = RP*(1.0 + e)/(1.0 - e)

dv_p = dv_escape(RP, a)
dv_a = dv_escape(ra, a)
dv_saving = (dv_a - dv_p)/dv_p*100.0          # what apoapsis costs extra

t_p = t_parabolic(RP, RSOI)                   # constant: q = RP for every e
t_a = t_parabolic(ra, RSOI)
half_period = np.pi*np.sqrt(a**3/MU)          # coast from apoapsis to periapsis

t_pen_burn = (t_p - t_a)/t_a*100.0            # clock started at the burn
t_pen_wait = (t_p + half_period - t_a)/t_a*100.0   # clock started at apoapsis

# ---------------------------------------------------------------------------
#  Figure
# ---------------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(11.6, 6.2))
ax2 = ax.twinx()

ax.axhline(0.0, color=INK, lw=1.2, alpha=0.6)

lt1, = ax.plot(e, t_pen_burn, color=BLUE, lw=2.8,
               label="time penalty, clock started at the burn")
lt2, = ax.plot(e, t_pen_wait, color=GREEN, lw=2.6, ls=(0, (5, 3)),
               label="time penalty, clock started at apoapsis (includes the coast down)")
ld,  = ax2.plot(e, dv_saving, color=MAROON, lw=2.8,
                label=r"$\Delta v$ penalty of burning at apoapsis instead")

# where the coast-inclusive curve changes its mind
sign_change = np.where(np.diff(np.sign(t_pen_wait[1:])) != 0)[0] + 1
for i in sign_change:
    ax.plot(e[i], 0.0, "o", color=GREEN, ms=8, mfc=CREAM, mew=2.0, zorder=9)
    ax.annotate(fr"$e={e[i]:.3f}$", xy=(e[i], 0.0), xytext=(e[i], 3.6),
                color=GREEN, fontsize=10, ha="center",
                arrowprops=dict(arrowstyle="-", color=GREEN, lw=1.0, alpha=0.6))

ax.set_xlabel(r"eccentricity,  $e$   (periapsis fixed at LEO, $r_p=6678$ km)",
              fontsize=12)
ax.set_ylabel(r"time to reach $R_{\rm SOI}$, periapsis relative to apoapsis  [%]",
              fontsize=12)
ax2.set_ylabel(r"$\Delta v$ penalty of burning at apoapsis  [%]",
               fontsize=12, color=MAROON)
ax2.tick_params(axis="y", colors=MAROON)

ax.set_xlim(0, 0.95)
ax.set_ylim(-26, 8)
ax2.set_ylim(-260, 1000)

ax.grid(True, color=INK, alpha=0.18, lw=0.7)
ax.spines["top"].set_visible(False)
ax2.spines["top"].set_visible(False)
ax2.spines["right"].set_color(MAROON)

ax.text(0.012, 7.5,
        "above the line: the coast down to periapsis has eaten the saving",
        color=INK, fontsize=10, alpha=0.85, ha="left", va="top")
ax.text(0.90, -25.4,
        "below the line: escaping from periapsis reaches the SOI sooner",
        color=INK, fontsize=10, alpha=0.85, ha="right", va="bottom")

ax.legend(handles=[lt1, lt2, ld], frameon=False, fontsize=10.5,
          loc="lower left", bbox_to_anchor=(0.015, 0.05))

ax.set_title("Escaping from periapsis: always cheaper, and slower only at the extremes",
             fontsize=13.5)
fig.subplots_adjust(left=0.085, right=0.895, top=0.92, bottom=0.12)
fig.savefig("img/fig-man-escape-penalty.png", dpi=160, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-man-escape-penalty.png")


# ---------------------------------------------------------------------------
#  Verification
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    print(f"\nR_SOI = {RSOI:,.0f} km   (dt/dq changes sign at R/2 = {RSOI/2:,.0f} km)")

    print("\nclosed form against Barker's equation")
    for q in (RP, 20034.0, 126885.0, 400000.0):
        D = np.sqrt(RSOI/q - 1.0)
        barker = 0.5*np.sqrt((2*q)**3/MU)*(D + D**3/3.0)
        print(f"  q={q:10.1f} km:  closed={t_parabolic(q, RSOI):13.3f} s   "
              f"Barker={barker:13.3f} s   diff={abs(barker-t_parabolic(q,RSOI)):.2e}")

    print("\nthe apoapsis of the whole family stays inside R_SOI/2, so dt/dq > 0 throughout")
    print(f"  largest apoapsis, e={e[-1]:.2f}:  {ra[-1]:,.0f} km  "
          f"({'inside' if ra[-1] < RSOI/2 else 'OUTSIDE'})")

    print(f"\n{'e':>5} {'r_a [km]':>11} {'dv_p':>7} {'dv_a':>7} {'dv pen %':>9} "
          f"{'t_p [d]':>8} {'t_a [d]':>8} {'t pen %':>8} {'t pen+coast %':>14}")
    for ev in (0.0, 0.1, 0.3, 0.5, 0.7, 0.9, 0.95):
        i = int(np.argmin(np.abs(e - ev)))
        print(f"{e[i]:5.2f} {ra[i]:11,.0f} {dv_p[i]:7.4f} {dv_a[i]:7.4f} "
              f"{dv_saving[i]:9.1f} {t_p/86400:8.3f} {t_a[i]/86400:8.3f} "
              f"{t_pen_burn[i]:8.2f} {t_pen_wait[i]:14.2f}")

    assert np.all(t_pen_burn <= 1e-9), "from the burn, periapsis is never slower"
    print("\nfrom the burn: penalty <= 0 everywhere, so periapsis is never slower")

    xs = e[1:][np.where(np.diff(np.sign(t_pen_wait[1:])) != 0)[0] + 1]
    print("from apoapsis: the coast flips the sign at e = "
          + ", ".join(f"{v:.3f}" for v in xs))
    print(f"  worst for periapsis: {t_pen_wait.max():+.1f}% at e={e[np.argmax(t_pen_wait)]:.3f}")
    print(f"  best  for periapsis: {t_pen_wait.min():+.1f}% at e={e[np.argmin(t_pen_wait)]:.3f}")
    assert len(xs) == 2, "expected the coast-inclusive curve to change sign twice"

    print(f"\ndv is minimised at periapsis for every e: max penalty "
          f"{dv_saving.max():.0f}% at e={e[np.argmax(dv_saving)]:.2f}")
    assert np.all(dv_saving >= -1e-9)