"""Apse-line rotation: one radial impulse at a crossing point, and when to prefer it."""
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.40
p = a*(1 - e**2)
eta = np.radians(60.0)                       # rotation, exaggerated for the picture

figL, axL = plt.subplots(figsize=(7.6, 6.4))
figR, axR = plt.subplots(figsize=(9.0, 5.8))

# ---- panel A : the two orbits and the crossings ---------------------------
f = np.linspace(0, 2*np.pi, 900)
r = p/(1 + e*np.cos(f))
axL.plot(r*np.cos(f), r*np.sin(f), color=BLUE, lw=2.2, label="initial orbit")
axL.plot(r*np.cos(f + eta), r*np.sin(f + eta), color=GREEN, lw=2.2, label="final orbit")

for fc, col in ((eta/2, MAROON), (eta/2 + np.pi, PURPLE)):
    rc = p/(1 + e*np.cos(fc))
    x, y = rc*np.cos(fc), rc*np.sin(fc)
    axL.plot(x, y, "o", color=col, ms=11, zorder=6)
    u = np.array([np.cos(fc), np.sin(fc)])   # radial direction
    vr = np.sqrt(mu/p)*e*np.sin(fc)
    L = 3400*np.sign(-vr)
    axL.annotate("", xy=(x + L*u[0], y + L*u[1]), xytext=(x, y),
                 arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=2.6))
axL.text(9000, 5200, r"$\Delta \mathbf{v}$ is purely" "\n" r"$\bf radial$",
         color=AMBER, fontsize=12)

# apse lines
axL.plot([0, 1.35*p], [0, 0], color=BLUE, lw=1.0, ls=":")
axL.plot([0, 1.35*p*np.cos(eta)], [0, 1.35*p*np.sin(eta)], color=GREEN, lw=1.0, ls=":")
arc = np.linspace(0, eta, 60)
axL.plot(5200*np.cos(arc), 5200*np.sin(arc), color=INK, lw=1.3)
axL.text(5900*np.cos(eta/2), 5900*np.sin(eta/2), r"$\eta$", color=INK, fontsize=14)
axL.plot(0, 0, "o", color=INK, ms=8)
axL.text(0, -19500, "the crossings sit at  $f=\\eta/2$  and  $f=\\eta/2+\\pi$;\n"
                    "both cost exactly the same",
         color=INK, fontsize=10, ha="center", alpha=0.8)
axL.set_aspect("equal"); axL.axis("off")
axL.legend(frameon=False, fontsize=10.5, loc="upper left")
axL.set_title(r"Rotating the apse line by $\eta$ (same $p$, same $e$)", fontsize=12)

# ---- panel B : one impulse versus the apoapsis route ----------------------
et = np.radians(np.linspace(0, 180, 400))
single = 2*np.sqrt(mu/p)*e*np.sin(et/2)
ra = a*(1 + e)
route = 2*(np.sqrt(mu/ra) - np.sqrt(mu*(2/ra - 1/a)))
axR.plot(np.degrees(et), single, color=MAROON, lw=2.8,
         label=r"single impulse  $2\sqrt{\mu/p}\,e\sin(\eta/2)$")
axR.axhline(route, color=GREEN, lw=2.2,
            label="circularise at apoapsis, then re-eccentricise")
cross = 2*np.arcsin(route/(2*np.sqrt(mu/p)*e))
axR.plot(np.degrees(cross), route, "o", color=INK, ms=9)
axR.annotate(f"crossover  $\\eta$ = {np.degrees(cross):.1f}°",
             xy=(np.degrees(cross), route), xytext=(np.degrees(cross) + 12, route - 0.9),
             color=INK, fontsize=11,
             arrowprops=dict(arrowstyle="-|>", color=INK, lw=1.3))
axR.fill_between(np.degrees(et), 0, np.minimum(single, route), color=AMBER, alpha=0.13)
axR.set_xlabel(r"apse-line rotation  $\eta$  [deg]")
axR.set_ylabel(r"$\Delta v$   [km/s]")
axR.set_xlim(0, 180); axR.set_ylim(0, 5.4)
axR.legend(frameon=False, fontsize=10, loc="upper left")
axR.set_title(r"$a=12000$ km, $e=0.4$: cheapest route depends on $\eta$", fontsize=12)
despine(axR)

for f, name in ((figL, "fig-man-apse-rotation-geom"), (figR, "fig-man-apse-rotation-cost")):
    f.tight_layout()
    f.savefig(f"img/{name}.png", dpi=160, bbox_inches="tight")
    plt.close(f)
    print(f"wrote img/{name}.png")