"""Hohmann transfer: geometry, and the delta-v split as a function of R = r2/r1."""
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": False,
})

figL, axL = plt.subplots(figsize=(7.8, 6.4))
figR, axR = plt.subplots(figsize=(8.8, 5.8))

# ---- panel A : geometry ---------------------------------------------------
r1, r2 = 1.0, 3.2
at = 0.5*(r1 + r2)
et = (r2 - r1)/(r2 + r1)
th = np.linspace(0, 2*np.pi, 600)

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)

# transfer ellipse: periapsis on +x, so r = p/(1+e cos f)
f = np.linspace(0, np.pi, 400)
rt = at*(1 - et**2)/(1 + et*np.cos(f))
axL.plot(rt*np.cos(f), rt*np.sin(f), color=MAROON, lw=2.6)
rt2 = at*(1 - et**2)/(1 + et*np.cos(-f))
axL.plot(rt2*np.cos(-f), rt2*np.sin(-f), color=MAROON, lw=1.0, ls=(0, (4, 4)), alpha=0.55)

axL.plot(0, 0, "o", color=INK, ms=9)
axL.text(0.10, -0.30, "focus", color=INK, fontsize=10)

# burn arrows (tangential, i.e. along +y at periapsis and -y at apoapsis)
axL.annotate("", xy=(r1, 0.92), xytext=(r1, 0.16),
             arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=2.6))
axL.text(r1 + 0.10, 0.60, r"$\Delta v_1$", color=AMBER, fontsize=14)
axL.annotate("", xy=(-r2, -0.92), xytext=(-r2, -0.16),
             arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=2.6))
axL.text(-r2 + 0.12, -0.66, r"$\Delta v_2$", color=AMBER, fontsize=14)

axL.plot([0, r1], [0, 0], color=BLUE, lw=1.1, ls=":")
axL.plot([0, -r2], [0, 0], color=GREEN, lw=1.1, ls=":")
axL.text(0.46, 0.12, r"$r_1$", color=BLUE, fontsize=13)
axL.text(-1.9, 0.12, r"$r_2$", color=GREEN, fontsize=13)
axL.text(1.55, 2.15, "transfer ellipse\n" + r"$a_t=\frac{r_1+r_2}{2}$",
         color=MAROON, fontsize=12, ha="center")
axL.text(0, r2 + 0.22, "half a revolution: the only\ncoast on the transfer orbit",
         color=INK, fontsize=10, ha="center", alpha=0.75)

axL.set_aspect("equal"); axL.axis("off")
axL.set_xlim(-r2 - 0.55, r2 + 0.55); axL.set_ylim(-r2 - 0.55, r2 + 0.95)
axL.set_title("Geometry: two tangential burns, 180° apart", fontsize=12)

# ---- panel B : normalised delta-v vs R ------------------------------------
R = np.linspace(1.0, 40.0, 1400)
dv1 = np.sqrt(2*R/(1 + R)) - 1
dv2 = (1/np.sqrt(R))*(1 - np.sqrt(2/(1 + R)))
tot = dv1 + dv2

axR.plot(R, dv1, color=BLUE, lw=2.0, label=r"$\Delta v_1/v_{c1}$  (departure)")
axR.plot(R, dv2, color=GREEN, lw=2.0, label=r"$\Delta v_2/v_{c1}$  (arrival)")
axR.plot(R, tot, color=MAROON, lw=2.8, label=r"total")

imax = int(np.argmax(tot))
axR.plot(R[imax], tot[imax], "o", color=MAROON, ms=7)
axR.annotate(f"maximum {tot[imax]:.4f} at $R$ = {R[imax]:.2f}",
             xy=(R[imax], tot[imax]), xytext=(R[imax] + 4.0, tot[imax] + 0.045),
             color=MAROON, fontsize=10,
             arrowprops=dict(arrowstyle="-", color=MAROON, lw=1.0))
axR.axhline(np.sqrt(2) - 1, color=INK, lw=1.0, ls=(0, (5, 4)), alpha=0.7)
axR.text(30.5, np.sqrt(2) - 1 + 0.012, r"$\sqrt{2}-1$  (escape burn)",
         color=INK, fontsize=10, alpha=0.8)

axR.axvline(11.94, color=PURPLE, lw=1.2, ls=":")
axR.text(11.94, 0.028, "  11.94", color=PURPLE, fontsize=10, rotation=90, va="bottom")

axR.set_xlabel(r"radius ratio  $R=r_2/r_1$")
axR.set_ylabel(r"$\Delta v / v_{c1}$")
axR.set_xlim(1, 40); axR.set_ylim(0, 0.60)
axR.legend(frameon=False, fontsize=10, loc="center right")
axR.set_title("Cost is not monotone in $R$: going further can cost less", fontsize=12)
for s in ("top", "right"):
    axR.spines[s].set_visible(False)

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