"""Bi-elliptic transfer: geometry, and where it beats the Hohmann transfer."""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine

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

# ---- panel A : geometry ---------------------------------------------------
r1, r2, rb = 1.0, 5.0, 11.0
th = np.linspace(0, 2*np.pi, 700)
axL.plot(r1*np.cos(th), r1*np.sin(th), color=BLUE, lw=1.8)
axL.plot(r2*np.cos(th), r2*np.sin(th), color=GREEN, lw=1.8)

def half(rp, ra, sign=+1, **kw):
    a, e = 0.5*(rp + ra), (ra - rp)/(ra + rp)
    f = np.linspace(0, np.pi, 400)
    r = a*(1 - e**2)/(1 + e*np.cos(f))
    axL.plot(r*np.cos(f), sign*r*np.sin(f), **kw)

half(r1, rb, +1, color=MAROON, lw=2.6)          # ellipse 1: r1 -> rb
half(r2, rb, -1, color=PURPLE, lw=2.6)          # ellipse 2: rb -> r2

for (x, y, dx, dy, lab, col) in [
        (r1, 0, 0, 0.95, r"$\Delta v_1$", AMBER),
        (-rb, 0, 0, -0.95, r"$\Delta v_2$", AMBER),
        (r2*np.cos(np.pi), 0, 0, 0, "", AMBER)]:
    if lab:
        axL.annotate("", xy=(x + dx, y + dy), xytext=(x, y + 0.18*np.sign(dy)),
                     arrowprops=dict(arrowstyle="-|>", color=col, lw=2.4))
        axL.text(x + 0.22, y + 0.55*np.sign(dy), lab, color=col, fontsize=13)
axL.annotate("", xy=(r2, -0.95), xytext=(r2, -0.18),
             arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=2.4))
axL.text(r2 + 0.22, -0.62, r"$\Delta v_3$", color=AMBER, fontsize=13)

axL.plot(0, 0, "o", color=INK, ms=8)
for x, col in ((r1, BLUE), (r2, GREEN), (-rb, MAROON)):
    axL.plot([0, x], [0, 0], color=col, lw=0.9, ls=":")
axL.text(0.30, -0.85, r"$r_1$", color=BLUE, fontsize=13)
axL.text(3.30, -0.95, r"$r_2$", color=GREEN, fontsize=13)
axL.text(-8.9, 0.42, r"$r_b \gg r_2$", color=MAROON, fontsize=13)
axL.text(-5.0, 3.75, "ellipse 1", color=MAROON, fontsize=12, ha="center")
axL.text(-3.0, -8.15, "ellipse 2", color=PURPLE, fontsize=12, ha="center")
axL.text(-2.5, 6.9, "burn 2 is cheap: the spacecraft is\nbarely moving out at $r_b$",
         color=INK, fontsize=10, ha="center", alpha=0.78)
axL.set_aspect("equal"); axL.axis("off")
axL.set_xlim(-rb - 1.6, r2 + 2.2); axL.set_ylim(-9.2, 8.4)
axL.set_title("Three burns, two ellipses, one long detour", fontsize=12)

# ---- panel B : cost comparison -------------------------------------------
def hoh(R):
    return np.sqrt(2*R/(1 + R)) - 1 + (1/np.sqrt(R))*(1 - np.sqrt(2/(1 + R)))

def bie(R, Rs):
    return ((np.sqrt(2*Rs/(1 + Rs)) - 1)
            + np.sqrt(2/Rs)*(np.sqrt(R/(R + Rs)) - np.sqrt(1/(1 + Rs)))
            + (np.sqrt(2*Rs/(R*(R + Rs))) - 1/np.sqrt(R)))

def stop(R, Rs):
    """Three burns with r_b < r_2: a *stopover*, not a bi-elliptic transfer.

    Ellipse 2 now runs from its periapsis r_b up to its apoapsis r_2, so the
    spacecraft arrives at r_2 slower than circular and the third burn is
    prograde.  Every magnitude is taken positive -- which is exactly what the
    signed bi-elliptic formula fails to do once R > R*.
    """
    d1 = np.sqrt(2*Rs/(1 + Rs)) - 1
    d2 = abs(np.sqrt(2*R/(Rs*(R + Rs))) - np.sqrt(2/(Rs*(1 + Rs))))
    d3 = abs(1/np.sqrt(R) - np.sqrt(2*Rs/(R*(R + Rs))))
    return d1 + d2 + d3


R = np.linspace(2, 60, 1200)
axR.plot(R, hoh(R), color=MAROON, lw=2.8, label="Hohmann (2 impulses)")
for Rs, col, ls in ((25, BLUE, "-"), (60, GREEN, "-"), (200, PURPLE, "-")):
    y = np.where(R < Rs, bie(R, np.minimum(R*0 + Rs, 1e9)), np.nan)
    axR.plot(R, y, color=col, lw=1.9, ls=ls, label=fr"bi-elliptic, $r_b/r_1={Rs}$")
axR.plot(R, (np.sqrt(2) - 1)*(1 + 1/np.sqrt(R)), color=INK, lw=1.6,
         ls=(0, (5, 3)), label=r"bi-parabolic ($r_b\to\infty$)")

Rs = np.linspace(25.0, 38.0, 400)
axR.plot(Rs, stop(Rs, 25.0), color=BLUE, lw=1.8, ls=(0, (3, 2.5)), alpha=0.75)
axR.annotate("", xy=(40.0, stop(40.0, 25.0)), xytext=(38.0, stop(38.0, 25.0)),
             arrowprops=dict(arrowstyle="-|>", color=BLUE, lw=1.8, alpha=0.75))
axR.text(41.2, 0.5665, "past $R=25$ this $r_b$ buys a\nstopover, not a bi-elliptic:\n"
         "the cost climbs away", color=BLUE, fontsize=9.5, alpha=0.9,
         ha="left", va="center")

axR.axvline(11.938765, color=AMBER, lw=1.3, ls=":")
axR.axvline(15.581715, color=AMBER, lw=1.3, ls=":")
axR.text(11.94, 0.335, "  11.94\n  bi-parabolic\n  first wins", color=AMBER, fontsize=9.5, va="bottom")
axR.text(15.58, 0.435, "   15.58\n   any $r_b>r_2$ wins", color=AMBER, fontsize=9.5, va="bottom")
axR.annotate("", xy=(15.581715, 0.428), xytext=(15.581715, 0.545),
             arrowprops=dict(arrowstyle="-", color=AMBER, lw=1.0))

# ---- panel B inset : the "merge" past R ~ 15.6 is not a merge --------------
#  Read at slide scale, the r_b/r_1 = 25 curve looks glued to the Hohmann curve
#  from about R = 15 onwards.  It is not: it runs strictly below, by up to
#  0.0030 v_c1 (0.56%) near R = 18.5, and it touches Hohmann only at R = 25 --
#  where r_b = r_2 and the "bi-elliptic" IS the Hohmann transfer.  It arrives
#  there with zero slope, which is the d(Dv_bi)/dR|_{R*=R} = 0 identity proved
#  in the appendix, and that is precisely why the two look welded together.
RB, RX, RG = 25.0, 14.098525, 18.501           # merge, crossover, widest gap

axZ = axR.inset_axes([0.345, 0.075, 0.365, 0.385])
Rz = np.linspace(13.2, 26.4, 2000)
axZ.plot(Rz, hoh(Rz), color=MAROON, lw=2.4)
axZ.plot(Rz[Rz <= RB], bie(Rz[Rz <= RB], RB), color=BLUE, lw=1.8)
axZ.plot(Rz[Rz >= RB], stop(Rz[Rz >= RB], RB), color=BLUE, lw=1.6,
         ls=(0, (3, 2.5)), alpha=0.75)
axZ.plot(RB, hoh(RB), "o", color=BLUE, ms=6, mfc=CREAM, mew=1.6, zorder=6)
axZ.axvline(RB, color=INK, lw=0.9, ls=":", alpha=0.7)
axZ.axvline(RX, color=AMBER, lw=1.0, ls=":")

axZ.annotate("", xy=(RG, hoh(RG)), xytext=(RG, bie(RG, RB)),
             arrowprops=dict(arrowstyle="<|-|>", color=INK, lw=1.1,
                             shrinkA=0, shrinkB=0, mutation_scale=9))
axZ.text(RG - 0.35, 0.5*(hoh(RG) + bie(RG, RB)),
         f"{hoh(RG) - bie(RG, RB):.4f}\n({100*(hoh(RG) - bie(RG, RB))/hoh(RG):.2f}%)",
         color=INK, fontsize=8.5, ha="right", va="center")
axZ.text(RX + 0.25, 0.52963, f"{RX:.2f}\nblue drops below", color=AMBER,
         fontsize=8.5, ha="left", va="bottom")
axZ.text(RB - 0.45, 0.52963, "they touch only at\n$R=r_b/r_1=25$", color=BLUE,
         fontsize=8.5, ha="right", va="bottom")

axZ.set_xlim(13.2, 26.4); axZ.set_ylim(0.5294, 0.5370)
axZ.set_xticks((14, 18, 22, 25)); axZ.set_yticks((0.530, 0.532, 0.534, 0.536))
axZ.tick_params(labelsize=8.5, length=3, pad=2)
axZ.grid(True, color=INK, alpha=0.15, lw=0.6)
axZ.set_title("zoom on the apparent merge: below, not on", fontsize=9, pad=3)
axR.indicate_inset_zoom(axZ, edgecolor=INK, alpha=0.55, lw=1.0)

# the same thing happens to every bi-elliptic curve at its own R = r_b/r_1
for Rs, col in ((25, BLUE), (60, GREEN)):
    axR.plot(Rs, hoh(Rs), "o", color=col, ms=5, mfc=CREAM, mew=1.4, zorder=6)

axR.set_xlabel(r"radius ratio  $R=r_2/r_1$")
axR.set_ylabel(r"$\Delta v_{\rm total}/v_{c1}$")
axR.set_xlim(2, 60); axR.set_ylim(0.30, 0.58)
axR.legend(frameon=False, fontsize=9.5, loc="lower right")
axR.set_title("Bi-elliptic pays only for a distant enough target", fontsize=12)
despine(axR)

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