"""Figures of the phasing-loop section of lunar-phasing-and-low-thrust.qmd.

fig-lpl-geometry   the worked example to scale: Earth, the three-loop phasing ellipse, the trans-lunar ellipse, the Moon at each perigee return;
                   and the clock: the Moon's angle from the perigee direction against time
fig-lpl-window     waiting time against loop apogee for k = 1..5 loops: which loops make which Moon phase possible

Run from the sfm-slides directory:   python3 codes/fig-lpl-phasing.py
"""
import numpy as np
import _figfont                                # all figure text at one scale
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _lunar import N_MOON, D, R_SOI, RE
from _lunarlt import *

MOONC = "#6b655b"
H = hohmann_to_moon()
phi0, phi_a, k = np.radians(123.068), np.radians(35.0), 3
W = required_wait(phi0, phi_a)
P3 = phasing_orbit(W/k)
T = P3["T"]
t1, f1 = soi_entry(phi0)

# ------------------------------------------------------------------ the geometry and the clock
fig, (ax, ay) = plt.subplots(1, 2, figsize=(14.0, 6.3), gridspec_kw=dict(width_ratios=[1.0, 0.95]))
f = np.linspace(0, 2*np.pi, 700)
rl = P3["a"]*(1 - P3["e"]**2)/(1 + P3["e"]*np.cos(f))
ax.plot(rl*np.cos(f)/1e3, rl*np.sin(f)/1e3, color=MAROON, lw=2.4, label=f"phasing loop, $T_{{ph}}={T/DAY:.3f}$ d, $r_a={P3['ra']/1e3:,.0f}$ thousand km")
ft = np.linspace(0, f1, 500)
rt = H["a"]*(1 - H["e"]**2)/(1 + H["e"]*np.cos(ft))
ax.plot(rt*np.cos(ft)/1e3, rt*np.sin(ft)/1e3, color=GREEN, lw=2.2, label=f"trans-lunar ellipse to the sphere of influence ({t1/HOUR:.1f} h)")
ftd = np.linspace(f1, np.pi, 200)
rtd = H["a"]*(1 - H["e"]**2)/(1 + H["e"]*np.cos(ftd))
ax.plot(rtd*np.cos(ftd)/1e3, rtd*np.sin(ftd)/1e3, color=GREEN, lw=1.2, ls=(0, (3, 3)), alpha=0.6)
th = np.linspace(0, 2*np.pi, 500)
ax.plot(D*np.cos(th)/1e3, D*np.sin(th)/1e3, color=MOONC, lw=1.0, ls=(0, (4, 4)), alpha=0.7)
ax.add_patch(Circle((0, 0), RE/1e3, color=BLUE, alpha=0.35))
ax.plot([0, 430], [0, 0], color=INK, lw=0.8, ls=":")
ax.text(432, 0, "perigee\ndirection", fontsize=9.5, va="center", ha="left", color=INK)
for j in range(k + 1):
    ang = phi_a + N_MOON*j*T
    pm = D*np.array([np.cos(ang), np.sin(ang)])/1e3
    ax.plot(*pm, "o", color=MOONC, ms=12, mfc=CREAM, mew=2.0, zorder=5)
    lab = "injection into the loop" if j == 0 else (f"{j} loop{'s' if j > 1 else ''} later" if j < k else "3 loops later: trans-lunar burn")
    off = np.array([np.cos(ang), np.sin(ang)])*(78 if j in (1, 2) else 70)
    ax.text(*(pm + off), f"{np.degrees(ang):.1f}$^\\circ$\n{lab}", fontsize=9, ha="center", va="center", color=MOONC)
    arc = np.linspace(0, ang, 60)
    ax.plot(70*np.cos(arc)*(1 + 0.12*j), 70*np.sin(arc)*(1 + 0.12*j), color=MOONC, lw=0.9)
# the Moon at arrival, with its sphere of influence
t_arr = W + t1
ang_arr = phi_a + N_MOON*t_arr
pm = D*np.array([np.cos(ang_arr), np.sin(ang_arr)])/1e3
ax.add_patch(Circle(pm, R_SOI/1e3, fill=False, color=MOONC, lw=1.2, ls=(0, (3, 3))))
ax.plot(*pm, "o", color=MOONC, ms=9, mfc=CREAM, mew=2.0, zorder=5)
ax.text(-300, -120, f"arrival: the craft reaches the\nsphere of influence ({np.degrees(ang_arr):.1f}$^\\circ$)", fontsize=9, ha="left", color=MOONC)
ax.plot(rt[-1]*0 + (H['a']*(1 - H['e']**2)/(1 + H['e']*np.cos(f1)))*np.cos(f1)/1e3, (H['a']*(1 - H['e']**2)/(1 + H['e']*np.cos(f1)))*np.sin(f1)/1e3, "o", color=INK, ms=7, zorder=6)
ax.set_aspect("equal"); ax.set_xlim(-480, 540); ax.set_ylim(-440, 520)
ax.set_xlabel("[1000 km], inertial frame, the perigee direction along $+x$")
ax.legend(frameon=False, fontsize=9.5, loc="lower left")
ax.set_title("Three phasing loops, to scale ($\\phi_a=35^\\circ$, $\\phi_0=123.07^\\circ$)", fontsize=12)
despine(ax)

tt = np.linspace(0, W, 400)/DAY
ay.plot(tt, np.degrees(phi_a + N_MOON*tt*DAY), color=MOONC, lw=2.6, label="the Moon, $\\phi(t)=\\phi_a+n_{\\rm Moon}t$")
ay.axhline(np.degrees(phi0), color=GREEN, lw=1.4, ls="--")
ay.text(0.15, np.degrees(phi0) + 1.5, "the angle the trans-lunar burn needs, $\\phi_0$", color=GREEN, fontsize=10)
for j in range(k + 1):
    tj = j*T/DAY
    ay.axvline(tj, color=MAROON, lw=1.0, ls=":")
    ay.plot(tj, np.degrees(phi_a + N_MOON*j*T), "o", color=MAROON, ms=9)
    ay.text(tj + 0.05, 24, "perigee\nburn 1" if j == 0 else ("perigee\nburn 2\n(trans-lunar)" if j == k else f"perigee\nreturn {j}"), fontsize=9, color=MAROON, va="bottom")
ay.annotate("", xy=(W/DAY, np.degrees(phi0)), xytext=(0, np.degrees(phi_a)), arrowprops=dict(arrowstyle="-", color="none"))
ay.text(3.55, 40, f"$W=(\\phi_0-\\phi_a)/n_{{\\rm Moon}}={W/DAY:.3f}$ d\n$=k\\,T_{{ph}}$ with $k=3$", fontsize=11, color=INK)
ay.set_xlim(-0.1, W/DAY + 0.9); ay.set_ylim(15, 135)
ay.set_xlabel("time since the first burn [days]"); ay.set_ylabel("the Moon's angle from the perigee direction [deg]")
ay.legend(frameon=False, fontsize=10, loc="upper left")
ay.set_title("The clock: the loops make the perigee arrive when the Moon is at $\\phi_0$", fontsize=12)
despine(ay)
fig.tight_layout()
fig.savefig("img/fig-lpl-geometry.png", dpi=150, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-lpl-geometry.png")

# ------------------------------------------------------------------ which loops make which Moon phase possible
cap = D - 1.5*R_SOI
month = 2*np.pi/N_MOON/DAY
fig, ax = plt.subplots(figsize=(10.4, 6.2))
Wd = np.linspace(0.02, month, 800)
cols = [BLUE, GREEN, MAROON, PURPLE, AMBER]
for kk, col in zip(range(1, 6), cols):
    T_leo = 2*np.pi*np.sqrt(R0**3/MU_E)
    wk = Wd[Wd*DAY/kk > 1.02*T_leo]                          # a loop must be longer than the parking orbit
    ra = np.array([phasing_orbit(w*DAY/kk)["ra"] for w in wk])/1e3
    ok = ra <= cap/1e3
    ax.plot(wk[ok], ra[ok], color=col, lw=2.6, label=f"$k={kk}$")
    ax.plot(wk[~ok], ra[~ok], color=col, lw=1.2, ls=(0, (3, 3)), alpha=0.5)
ax.axhline(cap/1e3, color=INK, lw=1.0, ls="--")
ax.text(0.3, cap/1e3 + 6, f"cap on the apogee: $r_{{\\rm Moon}}-1.5\\,r_{{\\rm SOI}}={cap/1e3:,.0f}$ thousand km", fontsize=10)
ax.axhline(D/1e3, color=MOONC, lw=1.0, ls=":")
ax.text(0.3, D/1e3 + 6, "the Moon's orbit", fontsize=10, color=MOONC)
ax.plot(W/DAY, P3["ra"]/1e3, "o", color=INK, ms=11, zorder=5)
ax.annotate(f"worked example: $W={W/DAY:.2f}$ d, $k=3$,\n$r_a={P3['ra']/1e3:,.0f}$ thousand km", xy=(W/DAY, P3["ra"]/1e3), xytext=(8.8, 70),
            fontsize=10.5, arrowprops=dict(arrowstyle="-|>", color=INK, lw=1.2))
ax.set_xlim(0, month); ax.set_ylim(0, 420)
ax.set_xlabel("waiting time $W$ between the first burn and the trans-lunar burn [days]   (one revolution of the Moon in the model: 27.28 d)")
ax.set_ylabel("apogee radius of each of the $k$ loops [1000 km]")
ax.legend(frameon=False, fontsize=11, loc="upper left", ncol=5, bbox_to_anchor=(0.0, 0.93))
ax.set_title("Every waiting time has a loop count: $k$ equal loops of period $W/k$", fontsize=12)
despine(ax)
fig.tight_layout()
fig.savefig("img/fig-lpl-window.png", dpi=150, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-lpl-window.png")