"""Radial, transverse and total speed against true anomaly, for a family of
eccentricities -- and the separate role played by the semi-major axis.

With  v_r = (mu/h) e sin f,  v_perp = (mu/h)(1 + e cos f)  and  mu/h = sqrt(mu/p),
dividing by sqrt(mu/a) leaves curves that depend on e alone:

    v_r / sqrt(mu/a)    = e sin f      / sqrt(1 - e^2)
    v_perp / sqrt(mu/a) = (1 + e cos f)/ sqrt(1 - e^2)
    v / sqrt(mu/a)      = sqrt(1 + 2e cos f + e^2) / sqrt(1 - e^2)

so e fixes the *shape* of every velocity curve and a only sets the scale.
"""

import os
import numpy as np
import matplotlib.pyplot as plt

os.makedirs("img", exist_ok=True)

CREAM, INK = "#fcfbf7", "#2f2a24"
MAROON, BLUE, GREEN, AMBER, PURPLE = "#6b1f1f", "#2f5ea8", "#2d6a2d", "#b45309", "#7e22ce"

plt.rcParams.update({
    "figure.facecolor": CREAM, "savefig.facecolor": CREAM, "axes.facecolor": CREAM,
    "axes.edgecolor": INK, "axes.labelcolor": INK, "text.color": INK,
    "xtick.color": INK, "ytick.color": INK, "font.size": 11,
    "axes.grid": True, "grid.alpha": 0.22, "grid.color": INK,
})

ecc = [0.0, 0.2, 0.4, 0.6, 0.8]
colors = [INK, BLUE, GREEN, AMBER, PURPLE]

f = np.linspace(0.0, 360.0, 2001)
fr = np.deg2rad(f)

fig, ax = plt.subplots(2, 2, figsize=(13.0, 8.4))
axR, axT, axV, axA = ax[0, 0], ax[0, 1], ax[1, 0], ax[1, 1]

for e, col in zip(ecc, colors):
    s = 1.0/np.sqrt(1.0 - e**2)                     # sqrt(a/p)
    vr = e*np.sin(fr)*s
    vt = (1.0 + e*np.cos(fr))*s
    v = np.sqrt(1.0 + 2*e*np.cos(fr) + e**2)*s
    kw = dict(color=col, lw=2.2 if e else 1.5,
              ls="-" if e else (0, (4, 3)), label=f"$e = {e:g}$")
    axR.plot(f, vr, **kw)
    axT.plot(f, vt, **kw)
    axV.plot(f, v, **kw)
    if e > 0.0:                                     # r = a  <=>  cos f = -e
        fa = np.rad2deg(np.arccos(-e))
        axV.plot([fa, 360.0 - fa], [1.0, 1.0], "o", color=col, ms=6, zorder=6)
        axR.plot([90.0, 270.0], [e*s, -e*s], "o", color=col, ms=5, zorder=6)

# ── radial component ────────────────────────────────────────────────────────
axR.axhline(0.0, color=INK, lw=0.9, alpha=0.5)
axR.set_title(r"radial:  $v_r = \sqrt{\mu/p}\;e\sin f$", fontsize=12.5, pad=8)
axR.set_ylabel(r"$v_r \,/\, \sqrt{\mu/a}$")
axR.text(70, -0.28, "outbound  $v_r > 0$", ha="center", fontsize=10, alpha=0.85)
axR.text(290, 0.22, "inbound  $v_r < 0$", ha="center", fontsize=10, alpha=0.85)
axR.text(180, -1.42, r"extrema at $f = 90^\circ, 270^\circ$ --- the ends of the"
         " latus rectum", ha="center", va="center", fontsize=9.5, alpha=0.85)
axR.set_ylim(-1.6, 1.6)

# ── transverse component ────────────────────────────────────────────────────
axT.set_title(r"transverse:  $v_\perp = \sqrt{\mu/p}\;(1 + e\cos f) = h/r$",
              fontsize=12.5, pad=8)
axT.set_ylabel(r"$v_\perp \,/\, \sqrt{\mu/a}$")
axT.annotate(r"periapsis: $\sqrt{\frac{1+e}{1-e}}$", xy=(0.0, 3.0),
             xytext=(48, 3.05), fontsize=10, va="center",
             arrowprops=dict(arrowstyle="->", color=INK, lw=1.0, alpha=0.8))
axT.annotate(r"apoapsis: $\sqrt{\frac{1-e}{1+e}}$", xy=(180.0, 0.333),
             xytext=(180, 1.55), fontsize=10, ha="center",
             arrowprops=dict(arrowstyle="->", color=INK, lw=1.0, alpha=0.8))
axT.set_ylim(0.0, 3.4)

# ── speed ───────────────────────────────────────────────────────────────────
axV.axhline(1.0, color=INK, lw=1.0, ls=(0, (5, 4)), alpha=0.7)
axV.text(8, 0.94, r"$v = \sqrt{\mu/a}$", fontsize=10, ha="left", va="top",
         alpha=0.9)
axV.set_title(r"speed:  $v = \sqrt{\mu/p}\;\sqrt{1 + 2e\cos f + e^{2}}$",
              fontsize=12.5, pad=8)
axV.set_ylabel(r"$v \,/\, \sqrt{\mu/a}$")
axV.set_xlabel("true anomaly $f$  [deg]")
axV.text(180, 3.10, r"markers: $r = a$ ($\cos f = -e$), where the speed equals"
         "\n" r"the circular speed $\sqrt{\mu/a}$ for that radius", ha="center",
         va="top", fontsize=9.5, alpha=0.85)
axV.set_ylim(0.0, 3.4)

for a_ in (axR, axT, axV):
    a_.set_xlim(0, 360)
    a_.set_xticks(range(0, 361, 90))
    for sp in ("top", "right"):
        a_.spines[sp].set_visible(False)

# one legend for the three normalised panels -- the family is the same in each
axR.legend(frameon=False, fontsize=9.5, ncol=3, handlelength=1.6,
           labelspacing=0.3, columnspacing=1.0, loc="upper right")

# ── what a does: the same curve, rescaled ───────────────────────────────────
MU_EARTH = 398600.4418                                # km^3/s^2
cases = [(7000.0, "LEO-like, $a = 7\\,000$ km", BLUE),
         (10000.0, "$a = 10\\,000$ km", GREEN),
         (26600.0, "GPS-like, $a = 26\\,600$ km", MAROON)]
labels = []
e_fix = 0.4
for a_km, lab, col in cases:
    p_km = a_km*(1.0 - e_fix**2)
    v_kms = np.sqrt(MU_EARTH/p_km)*np.sqrt(1.0 + 2*e_fix*np.cos(fr) + e_fix**2)
    axA.plot(f, v_kms, color=col, lw=2.2,
             label=rf"{lab}    $\sqrt{{\mu/a}} = {np.sqrt(MU_EARTH/a_km):.2f}$ km/s")

axA.set_title(rf"the role of $a$:  same shape, scaled by $\sqrt{{\mu/a}}$"
              rf"   ($e = {e_fix}$, Earth)", fontsize=12.5, pad=8)
axA.set_xlabel("true anomaly $f$  [deg]")
axA.set_ylabel("$v$  [km/s]")
axA.set_xlim(0, 360)
axA.set_xticks(range(0, 361, 90))
axA.set_ylim(0.0, 16.5)
axA.legend(frameon=False, fontsize=9.5, loc="upper left")
for sp in ("top", "right"):
    axA.spines[sp].set_visible(False)

plt.tight_layout()
fig.savefig("img/fig-velocity-components.png", dpi=160, bbox_inches="tight")
plt.close(fig)