"""Figures for lunar-trajectories-bates.qmd (img/fig-lun-flight.png, img/fig-lun-miss.png).

fig-lun-flight: (a) geocentric paths from a perigee injection at 320 km altitude, for several injection speeds, until they cross the Moon's
orbit (Moon on a circle of radius 384 400 km, lunar gravity neglected); (b) the flight time and the geocentric sweep angle against injection speed.
fig-lun-miss : how an injection error moves the point where the probe meets the Moon's orbit, measured against where the Moon is:
               the miss is D (d psi - n dt): the sweep angle and the Moon's travel during the (changed) flight time partly cancel.

Run from the sfm-slides directory:   python3 codes/fig-lun-flight.py
"""
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import brentq
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _lunar import *

D2R = np.pi / 180
r0 = RE + 320.0
vesc = np.sqrt(2 * MU_E / r0); vmin = v_min_to(r0, D)


def path(v0, n=400):
    eps, h, p, e = conic(r0, v0, 0.0)
    f1 = flight(r0, v0, 0.0, D)["f1"]
    f = np.linspace(0, f1, n)
    r = p / (1 + e * np.cos(f))
    return r * np.cos(f), r * np.sin(f)


fig, ax = plt.subplots(1, 2, figsize=(13.6, 5.8), gridspec_kw={"width_ratios": [1.05, 1]})
A = ax[0]
t_ = np.linspace(0, 2 * np.pi, 400)
A.plot(D * np.cos(t_), D * np.sin(t_), color=INK, lw=1.0, ls=":")
A.fill(RE / 1e3 * np.cos(t_), RE / 1e3 * np.sin(t_), color="#9bb7d8", zorder=5)
cols = [PURPLE, BLUE, GREEN, AMBER, MAROON]
for v0, c in zip((vmin, 10.85, 10.95, 11.1, 11.5), cols):
    x, y = path(v0); fl = flight(r0, v0, 0.0, D)
    lab = f"{v0:.2f} km/s: {fl['t']/HOUR:.0f} h, sweep {np.degrees(fl['psi']):.0f}°" + (" (minimum energy)" if v0 == vmin else "")
    A.plot(x / 1e3, y / 1e3, color=c, lw=2.0, label=lab)
    A.scatter([x[-1] / 1e3], [y[-1] / 1e3], color=c, s=26, zorder=6)
A.plot(D * np.cos(t_) / 1e3, D * np.sin(t_) / 1e3, color=INK, lw=1.0, ls=":")
A.scatter([0], [0], color=BLUE, s=40, zorder=7)
A.text(15, -35, "Earth", color=BLUE, fontsize=10.5)
A.text(-330, 340, "Moon's orbit", color=INK, fontsize=10.5)
A.set_aspect("equal"); A.set_xlim(-420, 420); A.set_ylim(-60, 420); A.set_xlabel("x [thousand km]"); A.set_ylabel("y [thousand km]")
A.legend(loc="upper center", fontsize=9, frameon=False, bbox_to_anchor=(0.5, -0.14), ncol=1)
A.set_title("(a)  injection at perigee, 320 km altitude, to the Moon's orbit", fontsize=11.5, loc="left"); despine(A)
B = ax[1]
vs = np.linspace(vmin + 1e-4, 11.3, 400)
fl = [flight(r0, v, 0.0, D) for v in vs]
B.plot(vs, [f["t"] / HOUR for f in fl], color=BLUE, lw=2.4); B.set_ylabel("flight time to the Moon's orbit [h]", color=BLUE); B.set_ylim(0, 130)
B.set_xlabel("injection speed [km/s]"); B.set_xlim(10.8, 11.3)
B2 = B.twinx(); B2.plot(vs, [np.degrees(f["psi"]) for f in fl], color=MAROON, lw=2.4); B2.set_ylabel("sweep angle  $\\psi_{ff}$ [deg]", color=MAROON); B2.set_ylim(100, 190)
B.axvline(vmin, color=PURPLE, lw=1.0, ls=":"); B.text(vmin - 0.003, 6, f"minimum energy\n{vmin:.3f} km/s", color=PURPLE, fontsize=9.5, ha="left")
B.axvline(vesc, color=INK, lw=1.0, ls=":"); B.text(vesc + 0.01, 30, f"escape\n{vesc:.3f} km/s", color=INK, fontsize=9.5)
v72 = brentq(lambda v: flight(r0, v, 0.0, D)["t"] / HOUR - 72, vmin + 1e-4, 11.3)
B.scatter([v72], [72], color=AMBER, s=45, zorder=6); B.text(v72 + 0.01, 76, f"72 h: {v72:.3f} km/s", color=AMBER, fontsize=10)
B.set_title("(b)  flight time and sweep angle against injection speed", fontsize=11.5, loc="left"); despine(B)
fig.tight_layout(); fig.savefig("img/fig-lun-flight.png", dpi=150, bbox_inches="tight")
print(f"vmin={vmin:.4f}, vesc={vesc:.4f}, speed for 72 h = {v72:.4f}; t(vmin)={flight(r0,vmin,0.0,D)['t']/HOUR:.2f} h")

# ------------------------------------------------ the miss
def S(v0, rr=r0, g0=0.0):
    fl = flight(rr, v0, g0, D); return fl["psi"] - N_MOON * fl["t"], fl["psi"], fl["t"]


vv = np.linspace(10.83, 11.3, 300)
h = 1e-5
dpsi = np.array([(S(v + h)[1] - S(v - h)[1]) / (2 * h) * D / 1e3 for v in vv])
dt = np.array([(S(v + h)[2] - S(v - h)[2]) / (2 * h) * N_MOON * D / 1e3 for v in vv])
net = dpsi - dt
vcancel = brentq(lambda v: S(v + h)[0] - S(v - h)[0], 10.84, 11.3)
dr = np.array([(S(v, r0 + 1e-2)[0] - S(v, r0 - 1e-2)[0]) / 2e-2 * D for v in vv])
dg = np.array([(S(v, r0, 1e-6)[0] - S(v, r0, -1e-6)[0]) / 2e-6 * D * D2R for v in vv])
rcancel = brentq(lambda v: S(v, r0 + 1e-2)[0] - S(v, r0 - 1e-2)[0], 10.84, 11.3)
fig, ax = plt.subplots(1, 2, figsize=(13.2, 5.0))
A = ax[0]
A.plot(vv, dpsi, color=MAROON, lw=2.0, label=r"sweep angle: $D\,\partial\psi_{ff}/\partial v_0$")
A.plot(vv, -dt, color=BLUE, lw=2.0, label=r"Moon's travel: $-D\,n\,\partial t/\partial v_0$")
A.plot(vv, net, color=INK, lw=2.8, label="net miss")
A.axhline(0, color=INK, lw=0.8); A.axvline(vcancel, color=PURPLE, lw=1.0, ls=":")
A.text(vcancel + 0.01, 160, f"net = 0 at {vcancel:.3f} km/s\n(sweep {np.degrees(S(vcancel)[1]):.0f}°)", color=PURPLE, fontsize=10)
A.set_ylim(-1000, 400); A.set_xlim(10.83, 11.3); A.set_xlabel("injection speed [km/s]"); A.set_ylabel("miss per 1 m/s of injection speed [km]")
A.legend(loc="lower right", fontsize=9.5, frameon=False); A.set_title("(a)  speed error", fontsize=11.5, loc="left"); despine(A)
B = ax[1]
B.plot(vv, dr, color=GREEN, lw=2.4, label="per 1 km of injection radius [km/km]")
B.plot(vv, dg / 100, color=AMBER, lw=2.4, label="per 0.01° of flight-path angle [km]")
B.axhline(0, color=INK, lw=0.8); B.axvline(rcancel, color=PURPLE, lw=1.0, ls=":")
B.text(rcancel + 0.01, 350, f"altitude error cancels at {rcancel:.3f} km/s", color=PURPLE, fontsize=10)
B.set_ylim(-250, 700); B.set_xlim(10.83, 11.3); B.set_xlabel("injection speed [km/s]"); B.set_ylabel("miss [km]")
B.legend(loc="center right", fontsize=9.5, frameon=False); B.set_title("(b)  altitude and flight-path-angle errors", fontsize=11.5, loc="left"); despine(B)
fig.tight_layout(); fig.savefig("img/fig-lun-miss.png", dpi=150, bbox_inches="tight")
print(f"speed errors cancel at v0={vcancel:.4f} (sweep {np.degrees(S(vcancel)[1]):.2f} deg, t={S(vcancel)[2]/HOUR:.2f} h); altitude errors at {rcancel:.4f}")
i11 = np.argmin(abs(vv - 11.0)); print(f"at 11.0: speed {net[i11]:.1f} km per m/s, altitude {dr[i11]:.1f} km/km, flight path {dg[i11]:.0f} km/deg")