"""Q9: LEO to a circular orbit at the Moon's distance -- Hohmann against a bi-elliptic transfer.

Left: the orbits to scale.  Right: total Delta v and flight time against the intermediate
apoapsis r_b (log axis).  r_b = 1.0e6 km is marked; 1.5e6 km is the Sun-Earth L1/L2 distance
(about the edge of the Earth's Hill sphere).

Run from the sfm-slides directory:   python3 codes/fig-mq-bielliptic-moon.py
"""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _mqcommon import mu, Re, vis, conic_xy

r1, r2 = Re + 300.0, 384400.0

def hohmann():
    at = (r1 + r2)/2
    d1 = vis(r1, at) - np.sqrt(mu/r1)
    d2 = np.sqrt(mu/r2) - vis(r2, at)
    return d1 + d2, np.pi*np.sqrt(at**3/mu)/86400

def bielliptic(rb):
    a1, a2 = (r1 + rb)/2, (r2 + rb)/2
    b1 = vis(r1, a1) - np.sqrt(mu/r1)
    b2 = vis(rb, a2) - vis(rb, a1)
    b3 = abs(np.sqrt(mu/r2) - vis(r2, a2))
    return b1 + b2 + b3, np.pi*(np.sqrt(a1**3/mu) + np.sqrt(a2**3/mu))/86400

dvH, tH = hohmann()
fig, (axL, axR) = plt.subplots(1, 2, figsize=(14.4, 5.8), gridspec_kw=dict(width_ratios=[1.15, 1]))

# ---- A : the orbits ---------------------------------------------------------------
rb = 1.0e6
t = np.linspace(0, 2*np.pi, 600)
axL.plot(r2*np.cos(t), r2*np.sin(t), color=INK, lw=1.2, ls=(0, (5, 4)), alpha=0.8)
axL.text(0, r2 + 25000, "target: circular, 384 400 km", ha="center", fontsize=10, color=INK)
at = (r1 + r2)/2
x, y = conic_xy(at, (r2 - r1)/(r2 + r1))
axL.plot(x, y, color=BLUE, lw=2.2, label=r"Hohmann: $\Delta v$ = 3.94 km/s, 5.0 days")
a1, a2 = (r1 + rb)/2, (r2 + rb)/2
x, y = conic_xy(a1, (rb - r1)/(rb + r1))
axL.plot(x, y, color=MAROON, lw=2.0, label=r"bi-elliptic burn 1 $\to$ apoapsis $r_b=10^{6}$ km")
x, y = conic_xy(a2, (rb - r2)/(rb + r2))
axL.plot(x, y, color=GREEN, lw=2.0, label=r"burn 2 at $r_b$ raises the periapsis to 384 400 km")
axL.plot(0, 0, "o", color=INK, ms=6)
axL.plot(-rb, 0, "o", color=MAROON, ms=8); axL.text(-rb + 20000, 40000, "$r_b$", color=MAROON, fontsize=12)
axL.plot(r2, 0, "o", color=GREEN, ms=8)
axL.text(r2 + 20000, -70000, "burn 3\n(retrograde)", color=GREEN, fontsize=10)
axL.set_aspect("equal"); axL.axis("off")
axL.set_xlim(-1.12e6, 6.0e5); axL.set_ylim(-6.6e5, 6.6e5)
axL.legend(frameon=False, fontsize=10, loc="upper center", bbox_to_anchor=(0.5, 0.0))
axL.set_title("To scale: the Earth is the dot at the focus", fontsize=11.5)

# ---- B : cost and time against r_b --------------------------------------------------
rbs = np.logspace(np.log10(4.5e5), np.log10(6e6), 200)
dv = np.array([bielliptic(x)[0] for x in rbs])
tt = np.array([bielliptic(x)[1] for x in rbs])
axR.semilogx(rbs, dv, color=MAROON, lw=2.6, label="bi-elliptic")
axR.axhline(dvH, color=BLUE, lw=2.2, label="Hohmann  (5.0 days)")
axR.set_xlabel(r"intermediate apoapsis $r_b$  [km]")
axR.set_ylabel(r"total $\Delta v$  [km/s]")
axR.set_ylim(3.5, 4.1)
for rb_, lab in ((1.0e6, r"$r_b=10^{6}$ km"), (1.5e6, "$r_b=1.5\\times10^{6}$ km\n(Sun--Earth L1/L2)")):
    d_, t_ = bielliptic(rb_)
    axR.plot(rb_, d_, "o", color=AMBER, ms=9, zorder=6)
    axR.annotate(f"{lab}\n{d_:.3f} km/s, saves {1000*(dvH - d_):.0f} m/s ({100*(dvH - d_)/dvH:.1f} %)\n{t_:.0f} days",
                 (rb_, d_), ((4.7e5 if rb_ < 1.2e6 else 2.0e6), (3.56 if rb_ < 1.2e6 else 3.80)), fontsize=10, color=AMBER,
                 arrowprops=dict(arrowstyle="-", color=AMBER, lw=0.8))
ax2 = axR.twinx()
ax2.semilogx(rbs, tt, color=GREEN, lw=1.8, ls=(0, (5, 4)), label="flight time (right axis)")
ax2.set_ylabel("flight time  [days]", color=GREEN)
ax2.set_ylim(0, 600)
h1, l1 = axR.get_legend_handles_labels(); h2, l2 = ax2.get_legend_handles_labels()
axR.legend(h1 + h2, l1 + l2, frameon=False, fontsize=10.5, loc="upper left")
axR.set_title(r"$R=r_2/r_1=57.6\ >\ 15.58$: any $r_b>r_2$ beats Hohmann --- by how much, and at what price?", fontsize=11)
for ax in (axR, ax2):
    ax.spines["top"].set_visible(False)

plt.tight_layout()
fig.savefig("img/fig-mq-bielliptic-moon.png", dpi=160, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-mq-bielliptic-moon.png")