"""Where on an ellipse escape is cheapest, and what a bad choice costs.

Minimum-energy escape is a parabola, so at radius r the burn must raise the
speed to sqrt(2 mu / r).  Burning prograde keeps the direction, so

    dv(f) = sqrt(2 mu / r) - sqrt(mu (2/r - 1/a)),      r = a(1-e^2)/(1+e cos f)

which is a strictly decreasing function of a/r: the cheapest point is the
smallest radius, i.e. periapsis.  This figure shows how much the choice
matters, which is not obvious -- for a near-circular orbit almost nothing, for
an eccentric one a factor of several.

All orbits share the same periapsis (LEO, 300 km altitude), so the family is
"same parking perigee, progressively more eccentric".

Run from the sfm-slides directory:

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

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


def dv_escape(f_deg, e, rp=RP, mu=MU):
    """Prograde impulse needed to reach escape speed, at true anomaly f."""
    a = rp/(1.0 - e)
    r = a*(1.0 - e*e)/(1.0 + e*np.cos(np.radians(f_deg)))
    return np.sqrt(2.0*mu/r) - np.sqrt(mu*(2.0/r - 1.0/a))


CASES = [(0.0, INK), (0.1, BLUE), (0.3, GREEN), (0.5, PURPLE),
         (0.7, AMBER), (0.9, MAROON)]

f = np.linspace(0.0, 180.0, 721)

fig, (axL, axR) = plt.subplots(1, 2, figsize=(13.4, 5.6))

for e, col in CASES:
    dv = dv_escape(f, e)
    lw = 2.6 if e == 0.0 else 2.2
    ls = (0, (4, 3)) if e == 0.0 else "-"
    axL.plot(f, dv, color=col, lw=lw, ls=ls, label=fr"$e={e:g}$")
    axR.plot(f, (dv/dv[0] - 1.0)*100.0, color=col, lw=lw, ls=ls,
             label=fr"$e={e:g}$")
    # ring the optimum, which is the left-hand edge for every curve
    axL.plot(0.0, dv[0], "o", color=col, ms=6, mfc=CREAM, mew=1.6, zorder=8)

for ax in (axL, axR):
    ax.set_xlim(0, 180)
    ax.set_xticks(np.arange(0, 181, 30))
    ax.set_xlabel(r"true anomaly of the burn,  $f$  [deg]", fontsize=11.5)
    ax.grid(True, color=INK, alpha=0.18, lw=0.7)
    despine(ax)

# name the two ends of the x axis rather than labelling inside the plot, where
# every horizontal strip is already occupied by a curve
# (left panel only: on the right the e=0.9 curve runs into the top corner)
axL.text(2.5, 3.45, "periapsis", color=INK, fontsize=9.5, alpha=0.75,
         rotation=90, ha="left", va="top")
axL.text(177.5, 3.45, "apoapsis", color=INK, fontsize=9.5, alpha=0.75,
         rotation=90, ha="right", va="top")

axL.set_ylabel(r"$\Delta v_{\rm escape}$  [km/s]", fontsize=11.5)
axL.set_title("What escape costs", fontsize=13)
axL.set_ylim(0, 3.5)

axR.set_ylabel("penalty over the periapsis burn  [%]", fontsize=11.5)
axR.set_title("What a non-optimal point costs", fontsize=13)
axR.set_ylim(0, 650)

# one legend, in the only quadrant that is empty in either panel
axR.legend(frameon=False, fontsize=10.5, ncol=2, loc="upper left",
           bbox_to_anchor=(0.06, 0.98))
axR.text(11, 300, "$e=0$: every point is periapsis,\nso the penalty is identically zero",
         color=INK, fontsize=9.5, alpha=0.8, ha="left", va="center")

fig.suptitle("Escape from an ellipse: the burn point is worth more than the burn",
             fontsize=14)
fig.subplots_adjust(left=0.07, right=0.985, top=0.86, bottom=0.13, wspace=0.24)
fig.savefig("img/fig-man-escape-point.png", dpi=160, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-man-escape-point.png")


# --------------------------------------------------------------------------
#  Verification
# --------------------------------------------------------------------------
if __name__ == "__main__":
    print("\ndv is monotone in f on [0, 180] -- periapsis is the unique minimum")
    for e, _ in CASES[1:]:
        d = np.diff(dv_escape(f, e))
        print(f"  e={e:3.1f}:  min increment = {d.min():+.3e}  (all >= 0)")

    print("\ncircular case: dv/v_c = sqrt(2)-1 at every f")
    vc = np.sqrt(MU/RP)
    print(f"  computed {dv_escape(0.0, 0.0)/vc:.10f}   exact {np.sqrt(2)-1:.10f}")

    print("\ncost at periapsis and apoapsis, and the penalty between them")
    print(f"  {'e':>5} {'dv(f=0)':>9} {'dv(f=180)':>10} {'penalty %':>10}")
    for e, _ in CASES:
        p, a_ = dv_escape(0.0, e), dv_escape(180.0, e)
        print(f"  {e:5.2f} {p:9.4f} {a_:10.4f} {(a_/p - 1)*100:10.1f}")