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

Bate's worked example (our constants): injection at perigee, r0 = 6698.137 km (320 km altitude), v0 = 10.846 km/s, crossing the Moon's
sphere of influence at lambda_1 = 30 deg on the leading side of the Moon.
fig-lun-patch   : (a) the geocentric ellipse up to the patch point with the Moon at injection and at arrival (the phase angle), seen from the Earth;
                  (b) the selenocentric hyperbola from the patch point to perilune and out again, with the Moon to scale.
fig-lun-perilune: perilune altitude, the aiming angle, the sense of motion and the phase angle at departure against lambda_1 for the same injection.

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

D2R = np.pi / 180
r0, v0, lam1 = RE + 320.0, 10.846, 30 * D2R
R = patch(r0, v0, 0.0, lam1)
fig, ax = plt.subplots(1, 2, figsize=(13.6, 6.0), gridspec_kw={"width_ratios": [1.1, 1]})
A = ax[0]
th = np.linspace(0, 2 * np.pi, 500)
# the geocentric ellipse in the frame where the patch point is on the +x axis side: rotate so the Moon at ARRIVAL is on the x-axis
p, e = R["p0"], R["e0"]
f = np.linspace(0, R["f1"], 400)
rr = p / (1 + e * np.cos(f))
rot = -(R["f1"] - R["eta1"])                   # puts the Moon at arrival at angle 0
x, y = rr * np.cos(f + rot) / 1e3, rr * np.sin(f + rot) / 1e3
A.plot(D * np.cos(th) / 1e3, D * np.sin(th) / 1e3, color=INK, lw=0.9, ls=":")
A.plot(x, y, color=MAROON, lw=2.4)
A.fill(RE / 1e3 * np.cos(th), RE / 1e3 * np.sin(th), color="#9bb7d8", zorder=5)
mo_arr = np.array([D, 0]) / 1e3
gam = R["gamma_dep"]
mo_inj = D / 1e3 * np.array([np.cos(-gam + 0.0), np.sin(-gam)])      # Moon at injection: gamma_dep ahead of the injection point, i.e. at +gam from the injection direction
inj_dir = rot                                                          # direction of injection point at f = 0 is angle rot
mo_inj = D / 1e3 * np.array([np.cos(inj_dir + gam), np.sin(inj_dir + gam)])
for pos, nm, c in ((mo_inj, "Moon at injection", GREEN), (mo_arr, "Moon at arrival", AMBER)):
    A.scatter(*pos, s=110, color=c, zorder=8, edgecolor=INK)
A.text(mo_inj[0] - 25, mo_inj[1] + 22, "Moon at injection", color=GREEN, fontsize=10.5, ha="right")
A.text(mo_arr[0] - 10, mo_arr[1] - 40, "Moon at arrival", color=AMBER, fontsize=10.5, ha="right")
A.plot([0, mo_inj[0]], [0, mo_inj[1]], color=GREEN, lw=1.0, ls="--")
A.plot([0, x[0]], [0, y[0]], color=MAROON, lw=1.0, ls="--")
a0 = np.linspace(inj_dir, inj_dir + gam, 60)
A.plot(70 * np.cos(a0), 70 * np.sin(a0), color=PURPLE, lw=1.6)
A.text(80 * np.cos(inj_dir + gam / 2) + 5, 80 * np.sin(inj_dir + gam / 2) + 5, r"$\gamma_{dep}=%.1f^\circ$" % np.degrees(gam), color=PURPLE, fontsize=11)
A.add_patch(Circle(mo_arr, R_SOI / 1e3, fill=False, ec=BLUE, lw=1.2, ls="--", zorder=6))
A.text(250, 60, "sphere of\ninfluence", color=BLUE, fontsize=10)
A.scatter([R["P"][0] / 1e3 * np.cos(0) - 0], [R["P"][1] / 1e3], s=0)
A.scatter([x[-1]], [y[-1]], color=INK, s=24, zorder=9); A.text(x[-1] - 80, y[-1] + 15, "patch point", fontsize=10.5, ha="right")
A.text(-55, 25, "Earth", color=BLUE, fontsize=10.5)
A.set_aspect("equal"); A.set_xlim(-60, 440); A.set_ylim(-260, 220); A.set_xlabel("[thousand km]")
A.set_title("(a)  the geocentric phase (inertial frame)", fontsize=11.5, loc="left"); despine(A)
# (b) selenocentric hyperbola
B = ax[1]
eps2, h2, e2, p2 = R["eps2"], R["h2"], R["e2"], R["p2"]
a2 = MU_M / (2 * eps2)                          # |a| of the hyperbola (energy positive)
rho = R["P"] - np.array([D, 0.0]); v2 = R["v2vec"]
# orbital elements of the hyperbola in the Moon-centred frame: eccentricity vector
evec = ((v2 @ v2 - MU_M / np.linalg.norm(rho)) * rho - (rho @ v2) * v2) / MU_M
ecc2 = np.linalg.norm(evec); perilune_dir = evec / ecc2
sgn = np.sign(h2)
finf = np.arccos(-1 / ecc2)
fs = np.linspace(-finf * 0.995, finf * 0.995, 800)
rp_ = p2 / (1 + ecc2 * np.cos(fs))
phi_p = np.arctan2(perilune_dir[1], perilune_dir[0])
xs = rp_ * np.cos(phi_p + sgn * fs) / 1e3; ys = rp_ * np.sin(phi_p + sgn * fs) / 1e3
keep = rp_ <= R_SOI * 1.0
B.plot(xs[keep], ys[keep], color=MAROON, lw=2.4)
B.add_patch(Circle((0, 0), RM / 1e3, color="#bdbdbd", ec=INK, zorder=4)); B.add_patch(Circle((0, 0), R_SOI / 1e3, fill=False, ec=BLUE, lw=1.2, ls="--"))
B.scatter([rho[0] / 1e3], [rho[1] / 1e3], color=INK, s=30, zorder=9); B.text(rho[0] / 1e3 + 2, rho[1] / 1e3 + 2, "patch point", fontsize=10.5)
pp = R["rp"] * perilune_dir / 1e3
B.scatter(*pp, color=AMBER, s=40, zorder=9); B.text(pp[0] + 6, pp[1] - 9, f"perilune: {R['rp']-RM:.0f} km altitude", color=AMBER, fontsize=10.5)
B.annotate("", xy=(rho[0] / 1e3 + 14 * v2[0] / np.linalg.norm(v2), rho[1] / 1e3 + 14 * v2[1] / np.linalg.norm(v2)), xytext=(rho[0] / 1e3, rho[1] / 1e3), arrowprops=dict(arrowstyle="-|>", color=GREEN, lw=2))
B.text(rho[0] / 1e3 + 6, rho[1] / 1e3 - 9, r"$\vec{v}_2$", color=GREEN, fontsize=12)
B.text(-20, -62, "sphere of influence", color=BLUE, fontsize=10); B.text(-14, 5, "Moon", fontsize=10.5)
B.set_aspect("equal"); B.set_xlim(-70, 70); B.set_ylim(-70, 70); B.set_xlabel("[thousand km]")
B.set_title("(b)  the selenocentric phase (Moon-centred, inertial axes)", fontsize=11.5, loc="left"); despine(B)
fig.tight_layout(); fig.savefig("img/fig-lun-patch.png", dpi=150, bbox_inches="tight")
print({k: (round(v / D2R, 3) if k in ("eta1", "g1", "gamma_dep", "zeta2") else round(float(v), 4)) for k, v in R.items() if k not in ("P", "v1vec", "v2vec")})

# perilune against lambda_1
lams = np.linspace(2, 85, 400) * D2R
res = [patch(r0, v0, 0.0, l) for l in lams]
fig, ax = plt.subplots(1, 3, figsize=(15.6, 4.8))
alt = np.array([r["rp"] - RM for r in res]); hh = np.array([r["h2"] for r in res]); gd = np.array([np.degrees(r["gamma_dep"]) for r in res])
A = ax[0]
sign_change = np.where(np.diff(np.sign(hh)) != 0)[0]
A.plot(np.degrees(lams), alt / 1e3, color=MAROON, lw=2.4)
A.axhline(0, color=INK, lw=0.8)
A.set_ylim(-2, 12)
A.set_xlabel(r"$\lambda_1$ [deg]"); A.set_ylabel("perilune altitude [thousand km]")
for i in sign_change:
    A.axvline(np.degrees(lams[i]), color=PURPLE, lw=1.0, ls=":"); A.text(np.degrees(lams[i]) + 0.8, 8, r"$h_2=0$: aimed at the Moon's centre", color=PURPLE, fontsize=9.5, rotation=90, va="center")
A.set_title(r"(a)  perilune altitude (negative: it hits the Moon)", fontsize=11.5, loc="left"); despine(A)
B = ax[1]
B.plot(np.degrees(lams), np.degrees([r["zeta2"] for r in res]), color=BLUE, lw=2.4)
B.axhline(0, color=INK, lw=0.8); B.set_xlabel(r"$\lambda_1$ [deg]"); B.set_ylabel(r"aiming angle $\zeta_2$ [deg]")
B.set_title("(b)  where $\\vec{v}_2$ points relative to the Moon's centre", fontsize=11.5, loc="left"); despine(B)
C = ax[2]
C.plot(np.degrees(lams), gd, color=GREEN, lw=2.4); C.set_xlabel(r"$\lambda_1$ [deg]"); C.set_ylabel(r"phase angle at departure $\gamma_{dep}$ [deg]")
C.set_title("(c)  how far ahead the Moon must be at injection", fontsize=11.5, loc="left"); despine(C)
fig.tight_layout(); fig.savefig("img/fig-lun-perilune.png", dpi=150, bbox_inches="tight")
i30 = np.argmin(abs(lams - 30 * D2R)); print("lambda1 values where h2 changes sign:", [np.degrees(lams[i]) for i in sign_change])