"""Why Hohmann wins among two-impulse transfers: the feasible set in (h, eps).

Both impulses satisfy exactly

    dv_i^2 = 2 eps + 3 mu/r_i - 2 h sqrt(mu) / r_i^{3/2},

so the total cost increases with eps and decreases with h.  The transfer conic
must actually reach both radii, which forces eps >= h^2/(2 r_i^2) - mu/r_i for
i = 1, 2.  The optimum therefore sits on the upper envelope of those two
parabolas -- and the envelope's corner is exactly the Hohmann ellipse.
"""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine

mu = 398600.4418
r1, r2 = 6678.137, 42164.0

hH = np.sqrt(2*mu*r1*r2/(r1 + r2))          # Hohmann angular momentum
epsH = -mu/(r1 + r2)                        # Hohmann energy

def floor_eps(h, r):                        # eps needed to reach radius r
    return h**2/(2*r**2) - mu/r

def dvtot(h, eps):
    a = 2*eps + 3*mu/r1 - 2*h*np.sqrt(mu)/r1**1.5
    b = 2*eps + 3*mu/r2 - 2*h*np.sqrt(mu)/r2**1.5
    return np.sqrt(np.maximum(a, 0)) + np.sqrt(np.maximum(b, 0))

fig, (axL, axR) = plt.subplots(1, 2, figsize=(14.8, 5.9))

# ---- panel A : the feasible set and the cost contours ---------------------
hh = np.linspace(0.70*hH, 1.30*hH, 420)
ee = np.linspace(-1.9*abs(epsH), 1.1*abs(epsH), 420)
H, E = np.meshgrid(hh, ee)
feas = (E >= floor_eps(H, r1)) & (E >= floor_eps(H, r2))
Z = np.where(feas, dvtot(H, E), np.nan)

cf = axL.contourf(H/hH, E/abs(epsH), Z, levels=22, cmap="YlOrBr_r", alpha=0.85)
cs = axL.contour(H/hH, E/abs(epsH), Z, levels=12, colors=[INK], linewidths=0.6, alpha=0.5)
cb = fig.colorbar(cf, ax=axL, pad=0.02)
cb.set_label(r"$\Delta v_1+\Delta v_2$   [km/s]")
cb.outline.set_edgecolor(INK)

axL.plot(hh/hH, floor_eps(hh, r1)/abs(epsH), color=BLUE, lw=2.2,
         label=r"reach $r_1$:  $\varepsilon=\frac{h^2}{2r_1^2}-\frac{\mu}{r_1}$")
axL.plot(hh/hH, floor_eps(hh, r2)/abs(epsH), color=GREEN, lw=2.2,
         label=r"reach $r_2$:  $\varepsilon=\frac{h^2}{2r_2^2}-\frac{\mu}{r_2}$")
axL.plot(1.0, epsH/abs(epsH), "o", color=MAROON, ms=11, zorder=6)
axL.annotate("Hohmann\n(the corner)", xy=(1.0, epsH/abs(epsH)),
             xytext=(1.105, -0.05), color=MAROON, fontsize=11, ha="left",
             arrowprops=dict(arrowstyle="-|>", color=MAROON, lw=1.6))
axL.annotate("", xy=(0.80, 0.60), xytext=(1.02, 0.20),
             arrowprops=dict(arrowstyle="-|>", color=INK, lw=1.8))
axL.text(0.755, 0.66, "cost grows\nwith $\\varepsilon$,\nfalls with $h$",
         color=INK, fontsize=10, ha="left")

axL.set_xlabel(r"angular momentum   $h/h_{\rm H}$")
axL.set_ylabel(r"specific energy   $\varepsilon/|\varepsilon_{\rm H}|$")
axL.set_title("Feasible transfer conics, coloured by total $\\Delta v$", fontsize=12)
axL.legend(frameon=False, fontsize=9.5, loc="lower right")
axL.set_xlim(0.70, 1.30); axL.set_ylim(-1.9, 1.1)
despine(axL)

# ---- panel B : the cost along the feasible boundary ----------------------
hb = np.linspace(0.86*hH, 1.14*hH, 900)
eb = np.maximum(floor_eps(hb, r1), floor_eps(hb, r2))
axR.plot(hb/hH, dvtot(hb, eb), color=MAROON, lw=2.8)
axR.plot(1.0, dvtot(hH, epsH), "o", color=MAROON, ms=10)
axR.annotate(f"minimum  {dvtot(hH, epsH):.4f} km/s",
             xy=(1.0, dvtot(hH, epsH)), xytext=(1.030, dvtot(hH, epsH) + 0.26),
             color=MAROON, fontsize=11,
             arrowprops=dict(arrowstyle="-|>", color=MAROON, lw=1.4))
axR.axvline(1.0, color=INK, lw=0.9, ls=":")
axR.text(0.905, 4.62, r"$r_2$ is apoapsis" "\n" r"(2nd burn tangential)",
         color=GREEN, fontsize=10, ha="center")
axR.text(1.095, 4.62, r"$r_1$ is periapsis" "\n" r"(1st burn tangential)",
         color=BLUE, fontsize=10, ha="center")
axR.set_xlabel(r"angular momentum   $h/h_{\rm H}$")
axR.set_ylabel(r"$\Delta v_1+\Delta v_2$   [km/s]")
axR.set_title("Cost on that boundary: one minimum, at $h=h_{\\rm H}$", fontsize=12)
axR.set_xlim(0.86, 1.14); axR.set_ylim(3.80, 4.95)
despine(axR)

plt.tight_layout()
fig.savefig("img/fig-man-hohmann-proof.png", dpi=160, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-man-hohmann-proof.png")