"""Oberth: the same impulse buys more energy where the spacecraft is fast,
and buys the most when it is aligned with the velocity."""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine

mu = 398600.4418
a, e = 12000.0, 0.5
p = a*(1 - e**2)
dv = 0.1                                      # 100 m/s test impulse

fig, (axL, axR) = plt.subplots(1, 2, figsize=(14.4, 5.4))

# ---- panel A : where to burn ---------------------------------------------
f = np.linspace(0, 2*np.pi, 700)
r = p/(1 + e*np.cos(f))
v = np.sqrt(mu*(2/r - 1/a))
deps = v*dv + 0.5*dv**2                       # tangential impulse
axL.plot(np.degrees(f), deps, color=MAROON, lw=2.8)
axL.axhline(0.5*dv**2, color=INK, lw=1.4, ls=(0, (5, 4)))
axL.text(120, 0.5*dv**2 + 0.03, r"$\frac{1}{2}\Delta v^2$ — all you get if $v=0$",
         color=INK, fontsize=10)
for ang, lab, col, tx, ty in ((0, "periapsis", GREEN, 22, 0.78),
                              (180, "apoapsis", BLUE, 196, 0.40)):
    i = int(np.argmin(abs(np.degrees(f) - ang)))
    axL.plot(ang, deps[i], "o", color=col, ms=10)
    axL.text(tx, ty, f"{lab}:  $v$={v[i]:.3f} km/s\n"
             fr"$\Delta\varepsilon$={deps[i]:.4f} km$^2$/s$^2$",
             color=col, fontsize=10.5, va="bottom")
axL.set_xlabel(r"true anomaly at which the 100 m/s burn is applied  [deg]")
axL.set_ylabel(r"energy gain  $\Delta\varepsilon$   [km$^2$/s$^2$]")
axL.set_xlim(0, 360); axL.set_xticks(range(0, 361, 60))
axL.set_ylim(0, 1.24)
axL.set_title(r"$a=12000$ km, $e=0.5$: a factor of %.1f between the apsides"
              % (deps[0]/deps[int(len(f)/2)]), fontsize=12)
despine(axL)

# ---- panel B : which way to point ----------------------------------------
beta = np.linspace(0, np.pi, 400)
for r0, v0, col, lab in ((p/(1 + e), np.sqrt(mu*(2/(p/(1 + e)) - 1/a)), GREEN, "at periapsis"),
                         (p/(1 - e), np.sqrt(mu*(2/(p/(1 - e)) - 1/a)), BLUE, "at apoapsis")):
    axR.plot(np.degrees(beta), v0*dv*np.cos(beta) + 0.5*dv**2, color=col, lw=2.6,
             label=f"{lab}  ($v$ = {v0:.3f} km/s)")
axR.axhline(0, color=INK, lw=0.9)
axR.axvline(90, color=INK, lw=0.9, ls=":")
axR.text(26, -0.62, "a purely radial burn changes\n"
                    "the energy only at second order",
         color=INK, fontsize=10)
axR.set_xlabel(r"angle $\beta$ between $\Delta\mathbf{v}$ and $\mathbf{v}$   [deg]")
axR.set_ylabel(r"$\Delta\varepsilon = v\,\Delta v\cos\beta + \frac{1}{2}\Delta v^2$")
axR.set_xlim(0, 180); axR.set_xticks(range(0, 181, 30))
axR.legend(frameon=False, fontsize=10.5, loc="upper right")
axR.set_title("Tangential is best; retrograde is worst", fontsize=12)
despine(axR)

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