"""Coaxial ellipse-to-ellipse transfers: aligned and anti-aligned apsides.

Same two orbit shapes in both panels; only the orientation of the target's apse
line differs.  Aligned, the cheapest apsidal transfer runs periapsis-of-inner to
apoapsis-of-outer (Example 2).  Anti-aligned there are *two* candidate 180 deg
apsidal transfers -- periapsis-to-periapsis and apoapsis-to-apoapsis -- and they
trade cost against time (Example 3).
"""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine

mu = 398600.4418

a1, e1 = 12000.0, 0.30                     # inner orbit
a2, e2 = 42000.0, 0.50                     # outer orbit
rp1, ra1 = a1*(1 - e1), a1*(1 + e1)        # 8400, 15600
rp2, ra2 = a2*(1 - e2), a2*(1 + e2)        # 21000, 63000

vis = lambda r, a: np.sqrt(mu*(2/r - 1/a))


def transfer(rA, aFrom, rB, aTo):
    """A 180 deg apsidal transfer from radius rA on orbit `aFrom` to rB on `aTo`."""
    at = 0.5*(rA + rB)
    dv1 = abs(vis(rA, at) - vis(rA, aFrom))
    dv2 = abs(vis(rB, aTo) - vis(rB, at))
    return at, abs(rB - rA)/(rA + rB), dv1, dv2, dv1 + dv2, np.pi*np.sqrt(at**3/mu)/3600


def conic(a, e, rot=0.0, n=800):
    """Full conic about the focus, apse line rotated by `rot` from +x."""
    f = np.linspace(0, 2*np.pi, n)
    r = a*(1 - e**2)/(1 + e*np.cos(f))
    return r*np.cos(f + rot), r*np.sin(f + rot)


def motion_arrows(ax, x, y, fracs, color, size=17):
    """Small arrowheads along a curve, in the direction of travel."""
    n = len(x)
    for fr in fracs:
        i = int(fr*(n - 1)) % (n - 1)
        j = (i + 4) % (n - 1)
        ax.annotate("", xy=(x[j], y[j]), xytext=(x[i], y[i]), zorder=7,
                    arrowprops=dict(arrowstyle="-|>", color=color, lw=0,
                                    mutation_scale=size, shrinkA=0, shrinkB=0))


def draw_transfer(ax, rNear, rFar, upper, color, label, fracs=(0.30, 0.72),
                  second=False):
    """Half-ellipse joining rNear and rFar, 180 deg apart.

    `upper=True` puts the transfer's periapsis on +x, `False` on -x.  `second`
    selects the other half of the ellipse as the flown arc, which is what an
    inbound (outer -> inner) transfer traverses.  Motion is always
    counter-clockwise, i.e. prograde.
    """
    at = 0.5*(rNear + rFar)
    et = (rFar - rNear)/(rFar + rNear)
    rot = 0.0 if upper else np.pi
    f = np.linspace(np.pi, 2*np.pi, 400) if second else np.linspace(0, np.pi, 400)
    r = at*(1 - et**2)/(1 + et*np.cos(f))
    xs, ys = r*np.cos(f + rot), r*np.sin(f + rot)
    ax.plot(xs, ys, color=color, lw=3.0, zorder=5, label=label)
    motion_arrows(ax, xs, ys, fracs, color, size=19)


def burn(ax, x, up, lab, dx=2700, fs=12.0):
    """Prograde tangential burn arrow at an apsis on the x-axis."""
    ax.annotate("", xy=(x, 9900 if up else -9900), xytext=(x, 1700 if up else -1700),
                arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=2.8), zorder=9)
    ax.text(x + dx, 6100 if up else -6500, lab, color=AMBER,
            fontsize=fs, va="center", ha="left" if dx > 0 else "right")


def frame(ax, title):
    ax.set_aspect("equal"); ax.axis("off")
    ax.set_xlim(-68000, 68000); ax.set_ylim(-58000, 44000)
    ax.legend(frameon=False, fontsize=9.5, loc="upper left")
    ax.set_title(title, fontsize=13)


def orbits(ax, rot2):
    x1, y1 = conic(a1, e1)
    x2, y2 = conic(a2, e2, rot=rot2)
    ax.plot(x1, y1, color=BLUE, lw=2.0, label=f"orbit 1: $a={a1:.0f}$ km, $e={e1:.2f}$")
    ax.plot(x2, y2, color=GREEN, lw=2.0, label=f"orbit 2: $a={a2:.0f}$ km, $e={e2:.2f}$")
    motion_arrows(ax, x1, y1, (0.13, 0.63), BLUE)
    motion_arrows(ax, x2, y2, (0.16, 0.66), GREEN)
    ax.plot([-ra1, rp1], [0, 0], color=BLUE, lw=0.9, ls=":", zorder=2)
    ax.plot([rp2*np.cos(rot2), -ra2*np.cos(rot2)], [0, 0], color=GREEN, lw=0.9,
            ls=":", zorder=2)
    ax.plot(0, 0, "o", color=INK, ms=8, zorder=8)


# ============================ panel A : aligned ============================
# two candidate 180 deg apsidal transfers here as well
atA1, etA1, d1A1, d2A1, totA1, tA1 = transfer(rp1, a1, ra2, a2)   # peri -> apo (longer)
atA2, etA2, d1A2, d2A2, totA2, tA2 = transfer(ra1, a1, rp2, a2)   # apo -> peri (smaller)

figA, axA = plt.subplots(figsize=(8.6, 7.4))
orbits(axA, 0.0)

# A1: periapsis(1) on +x -> apoapsis(2) on -x; flown arc is the upper half
draw_transfer(axA, rp1, ra2, True, MAROON,
              f"A1  peri$\\to$apo: $a_t={atA1:.0f}$ km   {totA1:.3f} km/s, {tA1:.2f} h")
burn(axA, rp1, True, r"$\Delta v_1$")
burn(axA, -ra2, False, r"$\Delta v_2$")
axA.plot([rp1, -ra2], [0, 0], "o", color=MAROON, ms=9, zorder=9)

# A2: apoapsis(1) on -x -> periapsis(2) on +x; its periapsis is on -x, arc is lower
draw_transfer(axA, ra1, rp2, False, PURPLE,
              f"A2  apo$\\to$peri: $a_t={atA2:.0f}$ km   {totA2:.3f} km/s, {tA2:.2f} h",
              fracs=(0.26, 0.70))
burn(axA, -ra1, False, r"$\Delta v_1'$", dx=2700)
burn(axA, rp2, True, r"$\Delta v_2'$")
axA.plot([-ra1, rp2], [0, 0], "o", color=PURPLE, ms=9, zorder=9)

axA.text(0, -42000, f"A1  peri$\\to$apo:  ${d1A1:.3f}+{d2A1:.3f}={totA1:.3f}$ km/s "
                    f"in ${tA1:.3f}$ h", color=MAROON, fontsize=12, ha="center")
axA.text(0, -47500, f"A2  apo$\\to$peri:  ${d1A2:.3f}+{d2A2:.3f}={totA2:.3f}$ km/s "
                    f"in ${tA2:.3f}$ h", color=PURPLE, fontsize=12, ha="center")
axA.text(0, -53000, f"the longer route is cheaper by {1000*(totA2-totA1):.0f} m/s "
                    f"and slower by {tA1/tA2:.1f}$\\times$",
         color=INK, fontsize=12.5, ha="center", weight="bold")
frame(axA, r"Aligned apsides  ($\omega_2=\omega_1$)")

# ========================= panel B : anti-aligned ==========================
atS, etS, d1S, d2S, totS, tS = transfer(rp1, a1, rp2, a2)   # peri -> peri (smaller)
atL, etL, d1L, d2L, totL, tL = transfer(ra1, a1, ra2, a2)   # apo  -> apo  (longer)

figB, axB = plt.subplots(figsize=(8.6, 7.4))
orbits(axB, np.pi)

# T2: periapsis(1) on +x  ->  periapsis(2) on -x; flown arc is the upper half
draw_transfer(axB, rp1, rp2, True, MAROON,
              f"T2  peri$\\to$peri: $a_t={atS:.0f}$ km   {totS:.3f} km/s, {tS:.2f} h")
burn(axB, rp1, True, r"$\Delta v_1$")
burn(axB, -rp2, False, r"$\Delta v_2$", dx=-2700)
axB.plot([rp1, -rp2], [0, 0], "o", color=MAROON, ms=9, zorder=9)

# T1: apoapsis(1) on -x  ->  apoapsis(2) on +x; its periapsis is on -x, arc is lower
draw_transfer(axB, ra1, ra2, False, PURPLE,
              f"T1  apo$\\to$apo: $a_t={atL:.0f}$ km   {totL:.3f} km/s, {tL:.2f} h",
              fracs=(0.26, 0.68))
burn(axB, -ra1, False, r"$\Delta v_1'$", dx=2700)
burn(axB, ra2, True, r"$\Delta v_2'$")
axB.plot([-ra1, ra2], [0, 0], "o", color=PURPLE, ms=9, zorder=9)

axB.text(0, -42000, f"T2  peri$\\to$peri:  ${d1S:.3f}+{d2S:.3f}={totS:.3f}$ km/s "
                    f"in ${tS:.3f}$ h", color=MAROON, fontsize=12, ha="center")
axB.text(0, -47500, f"T1  apo$\\to$apo:  ${d1L:.3f}+{d2L:.3f}={totL:.3f}$ km/s "
                    f"in ${tL:.3f}$ h", color=PURPLE, fontsize=12, ha="center")
axB.text(0, -53000, f"the longer route is cheaper by {1000*(totS-totL):.0f} m/s "
                    f"and slower by {tL/tS:.1f}$\\times$",
         color=INK, fontsize=12.5, ha="center", weight="bold")
frame(axB, r"Anti-aligned apsides  ($\omega_2=\omega_1+180^\circ$)")

# ===================== panel C : outer -> inner (inbound) =====================
# Same two transfer ellipses as panel A, flown the other way round.  Time reversal
# leaves |dv| unchanged and turns every burn retrograde.
figC, axC = plt.subplots(figsize=(8.6, 7.4))
orbits(axC, 0.0)

# B1: apoapsis(2) on -x -> periapsis(1) on +x, i.e. A1's ellipse through the lower half
draw_transfer(axC, rp1, ra2, True, MAROON,
              f"B1  apo$\\to$peri: $a_t={atA1:.0f}$ km   {totA1:.3f} km/s, {tA1:.2f} h",
              fracs=(0.28, 0.70), second=True)
burn(axC, -ra2, True, r"$\Delta v_1$")
burn(axC, rp1, False, r"$\Delta v_2$")
axC.plot([-ra2, rp1], [0, 0], "o", color=MAROON, ms=9, zorder=9)

# B2: periapsis(2) on +x -> apoapsis(1) on -x, i.e. A2's ellipse through the upper half
draw_transfer(axC, ra1, rp2, False, PURPLE,
              f"B2  peri$\\to$apo: $a_t={atA2:.0f}$ km   {totA2:.3f} km/s, {tA2:.2f} h",
              fracs=(0.30, 0.72), second=True)
burn(axC, rp2, False, r"$\Delta v_1'$")
burn(axC, -ra1, True, r"$\Delta v_2'$", dx=2700)
axC.plot([rp2, -ra1], [0, 0], "o", color=PURPLE, ms=9, zorder=9)

axC.text(0, -42000, f"B1  apo$\\to$peri:  ${d2A1:.3f}+{d1A1:.3f}={totA1:.3f}$ km/s "
                    f"in ${tA1:.3f}$ h", color=MAROON, fontsize=12, ha="center")
axC.text(0, -47500, f"B2  peri$\\to$apo:  ${d2A2:.3f}+{d1A2:.3f}={totA2:.3f}$ km/s "
                    f"in ${tA2:.3f}$ h", color=PURPLE, fontsize=12, ha="center")
axC.text(0, -53000, "identical totals to Example 2 --- every burn is now retrograde",
         color=INK, fontsize=12.5, ha="center", weight="bold")
frame(axC, r"Outer $\to$ inner, aligned apsides")

# ============ panel D : all four legs on the two transfer ellipses ============
# A1 and B1 are the two halves of one ellipse; A2 and B2 the two halves of the
# other.  Solid = outbound (prograde), dashed = the return leg (retrograde).
# Each node therefore carries a single impulse magnitude, applied both ways.

def arc(rNear, rFar, upper, second, n=400):
    at = 0.5*(rNear + rFar)
    et = (rFar - rNear)/(rFar + rNear)
    rot = 0.0 if upper else np.pi
    f = np.linspace(np.pi, 2*np.pi, n) if second else np.linspace(0, np.pi, n)
    r = at*(1 - et**2)/(1 + et*np.cos(f))
    return r*np.cos(f + rot), r*np.sin(f + rot)


def leg(ax, rNear, rFar, upper, second, color, label, fracs=(0.32, 0.70)):
    xs, ys = arc(rNear, rFar, upper, second)
    ax.plot(xs, ys, color=color, lw=2.6 if second else 3.0, zorder=5, label=label,
            ls=(0, (7, 3)) if second else "-")
    motion_arrows(ax, xs, ys, fracs, color, size=19)


def node(ax, x, dv, dx, y):
    """One apsis: the same impulse, prograde outbound and retrograde inbound."""
    ax.annotate("", xy=(x, 10500), xytext=(x, -10500),
                arrowprops=dict(arrowstyle="<|-|>", color=AMBER, lw=2.6), zorder=9)
    ax.text(x + dx, y, f"{dv:.3f}", color=AMBER, fontsize=11.5,
            ha="left" if dx > 0 else "right", va="center")


figD, axD = plt.subplots(figsize=(8.8, 7.4))
orbits(axD, 0.0)

# ellipse 1 (a_t = 35700 km): periapsis 8400 on +x, apoapsis 63000 on -x
leg(axD, rp1, ra2, True,  False, MAROON,
    f"A1 out:  peri$\\to$apo, $a_t={atA1:.0f}$ km, {totA1:.3f} km/s")
leg(axD, rp1, ra2, True,  True,  MAROON, "B1 back: apo$\\to$peri, same ellipse")
# ellipse 2 (a_t = 18300 km): periapsis 15600 on -x, apoapsis 21000 on +x
leg(axD, ra1, rp2, False, False, PURPLE,
    f"A2 out:  apo$\\to$peri, $a_t={atA2:.0f}$ km, {totA2:.3f} km/s")
leg(axD, ra1, rp2, False, True,  PURPLE, "B2 back: peri$\\to$apo, same ellipse")

node(axD, -ra2, d2A1,  3200,  6000)
node(axD, -ra1, d1A2, -3200,  6000)
node(axD,  rp1, d1A1,  3200, -6000)
node(axD,  rp2, d2A2,  3200,  6000)
axD.plot([rp1, -ra2], [0, 0], "o", color=MAROON, ms=9, zorder=9)
axD.plot([-ra1, rp2], [0, 0], "o", color=PURPLE, ms=9, zorder=9)

axD.text(0, -42000, "two ellipses, four legs: each is flown out prograde and back "
                    "retrograde, for the same $\\Delta v$",
         color=INK, fontsize=11.5, ha="center")
axD.text(0, -47500, f"A1 = B1 = {totA1:.3f} km/s in {tA1:.3f} h        "
                    f"A2 = B2 = {totA2:.3f} km/s in {tA2:.3f} h",
         color=INK, fontsize=12.5, ha="center", weight="bold")
frame(axD, "All four legs, aligned apsides")
# six entries here, so give the legend a quiet backing to sit over the conics
axD.legend(frameon=True, fontsize=9.0, loc="upper left", facecolor=CREAM,
           edgecolor="none", framealpha=0.88)

for f_, name in ((figA, "fig-man-coaxial-aligned"),
                 (figC, "fig-man-coaxial-inbound"),
                 (figD, "fig-man-coaxial-all"),
                 (figB, "fig-man-coaxial-antialigned")):
    f_.tight_layout()
    f_.savefig(f"img/{name}.png", dpi=160, bbox_inches="tight")
    plt.close(f_)
    print(f"wrote img/{name}.png")

print(f"\naligned  A1 peri->apo : {d1A1:.5f}+{d2A1:.5f}={totA1:.5f} km/s  a_t={atA1:.1f}  e_t={etA1:.4f}  t={tA1:.3f} h")
print(f"aligned  A2 apo->peri : {d1A2:.5f}+{d2A2:.5f}={totA2:.5f} km/s  a_t={atA2:.1f}  e_t={etA2:.4f}  t={tA2:.3f} h")
print(f"anti T2  peri->peri: {d1S:.5f}+{d2S:.5f}={totS:.5f} km/s  a_t={atS:.1f}  e_t={etS:.4f}  t={tS:.3f} h")
print(f"anti T1  apo->apo  : {d1L:.5f}+{d2L:.5f}={totL:.5f} km/s  a_t={atL:.1f}  e_t={etL:.4f}  t={tL:.3f} h")
print(f"\nT1 cheaper by {totS-totL:.5f} km/s ({100*(totS-totL)/totL:.1f}%);  T2 faster by {tL-tS:.3f} h ({tL/tS:.2f}x)")