"""Verification of every number and identity quoted in lunar-phasing-and-low-thrust.qmd.

Part 1  the impulsive transfer the loops deliver (reproduced from patched-conics.qmd, Example 2a)
Part 2  the phasing loop: waiting time, loop period, the two burns, the window, and what the idealisation leaves out (numerical Earth+Moon run, J2)
Part 3  the spiral: symbolic derivations, a numerical integration, the eccentricity bound, where the model stops, propellant, power and time
Part 4  the worked examples and the comparison tables

Run from the sfm-slides directory:   python3 codes/lunar-phasing-lowthrust-verify.py        (about a minute)
"""
import numpy as np
import sympy as sp
from scipy.optimize import brentq
from _lunar import MU_E, MU_M, RE, RM, D, N_MOON, R_SOI
from _lunarlt import *
import _lunarlt as L

ok = lambda label, cond: print(f"  [{'ok' if cond else 'FAIL'}] {label}")
close = lambda x, y, tol: abs(x - y) <= tol
D2R, R2D = np.pi/180, 180/np.pi

# =========================================================================================== Part 1
print("Part 1: the transfer delivered by the loops (Example 2a of patched-conics.qmd)")
H = hohmann_to_moon()
ok(f"Hohmann ellipse to the Moon's orbit: a_H = {H['a']:.1f} km, e_H = {H['e']:.5f}, v_TLI = {H['v_p']:.4f} km/s, dv_TLI = {H['dv_tli']:.4f}",
   close(H["a"], 195539.1, 0.05) and close(H["e"], 0.96585, 5e-6) and close(H["v_p"], 10.8322, 5e-5) and close(H["dv_tli"], 3.1064, 5e-5))
ok(f"n_Moon = {N_MOON:.4e} rad/s = {N_MOON*R2D*DAY:.4f} deg/day (period {2*np.pi/N_MOON/DAY:.3f} d); the note's 2.6653e-6", close(N_MOON, 2.6653e-6, 5e-11))
ok(f"half period of the ellipse {H['T']/2/HOUR:.3f} h; the Moon turns {N_MOON*H['T']/2*R2D:.2f} deg meanwhile (65.71)", close(N_MOON*H["T"]/2*R2D, 65.71, 0.005))
# the design of Example 2a, found again from its helper: phase of the Moon at injection for a 100 km perilune
import _pc_transfer as pct
sols = pct.solve_aim(lambda d_: pct.hohmann_run(d_), 100.0, -30, 5)
phi0 = np.pi - N_MOON*(H["T"]/2 + sols[0])
Hh, te, pc = pct.hohmann_run(sols[0])
ok(f"Example 2a design: phi_0 = {phi0*R2D:.3f} deg, SOI entered {te/HOUR:.3f} h after injection, lambda_1 = {pc['lam1']*R2D:.2f}, perilune altitude {pc['rp'] - RM:.1f} km",
   close(phi0*R2D, 123.07, 0.005) and close(te/HOUR, 78.86, 0.005) and close(pc["lam1"]*R2D, 59.85, 0.005) and close(pc["rp"] - RM, 100.0, 0.1))
ok(f"  capture {pct.capture_dv(pc['eps2'], pc['rp']):.4f} km/s, total {H['dv_tli'] + pct.capture_dv(pc['eps2'], pc['rp']):.4f} km/s (0.7949, 3.9013)",
   close(pct.capture_dv(pc["eps2"], pc["rp"]), 0.7949, 5e-5) and close(H["dv_tli"] + pct.capture_dv(pc["eps2"], pc["rp"]), 3.9013, 5e-5))
ok(f"the helper soi_entry (used by the figures) agrees: {soi_entry(phi0)[0]/HOUR:.3f} h", close(soi_entry(phi0)[0]/HOUR, te/HOUR, 1e-4))
def alt_at_phi0(ph0):                                         # perilune altitude of Example 2a's patched-conic solution for a given phi_0
    return pct.hohmann_run((np.pi - ph0)/N_MOON - H["T"]/2)[2]["rp"] - RM


a_0, a_p, a_m = alt_at_phi0(phi0), alt_at_phi0(phi0 + D2R), alt_at_phi0(phi0 - D2R)
ok(f"perilune altitude {a_0:.1f} km at phi_0; {a_p:.0f} km at +1 deg and {a_m:.0f} km at -1 deg: {a_p - a_0:.0f} and {a_0 - a_m:.0f} km per degree of phi_0 (1 deg of the Moon = {1/(N_MOON*R2D*HOUR):.2f} h)",
   close(a_0, 100.0, 0.2) and close(a_p - a_0, 994, 3) and close(a_0 - a_m, 734, 3))
PHI0, T_SOI = phi0, te                                       # used below

# =========================================================================================== Part 2
print("Part 2: the phasing loop")
phi_a = 35.0*D2R
W = required_wait(PHI0, phi_a)
ok(f"waiting time W = (phi_0 - phi_a)/n_Moon = (123.068 - 35)/{N_MOON*R2D*DAY:.4f} = {W/DAY:.3f} d = {W/HOUR:.2f} h", close(W/DAY, 6.6747, 5e-4))
rows = {}
for k in (1, 2, 3, 4, 5):
    T = loop_period(W, k); P = phasing_orbit(T)
    rows[k] = P | dict(dv1=P["v_p"] - V_C0, dv2=H["v_p"] - P["v_p"])
ok("loops k = 1..5: T_ph = " + ", ".join(f"{rows[k]['T']/DAY:.3f}" for k in rows) + " d;  r_a = " + ", ".join(f"{rows[k]['ra']:,.0f}" for k in rows) + " km",
   all(close(rows[k]["T"]/DAY, v, 5e-4) for k, v in zip(rows, (6.6747, 3.3374, 2.2249, 1.6687, 1.3349)))
   and all(close(rows[k]["ra"], v, 1.0) for k, v in zip(rows, (292816, 181991, 137304, 112176, 95747))))
P3 = rows[3]
ok(f"k = 3: a_ph = {P3['a']:.1f} km, e_ph = {P3['e']:.5f}, r_a = {P3['ra']:.1f} km, v_p = {P3['v_p']:.4f} km/s", close(P3["a"], 71990.9, 0.1) and close(P3["e"], 0.90724, 5e-6) and close(P3["v_p"], 10.6695, 5e-5))
ok(f"  dv_1 = {P3['dv1']:.4f}, dv_2 = {P3['dv2']:.4f}, sum {P3['dv1'] + P3['dv2']:.4f} km/s = dv_TLI {H['dv_tli']:.4f}",
   close(P3["dv1"], 2.9437, 5e-5) and close(P3["dv2"], 0.1627, 5e-5) and close(P3["dv1"] + P3["dv2"], H["dv_tli"], 1e-12))
ok("the loop costs nothing in the impulsive model: dv_1 + dv_2 is the same for k = 1..5 and for any T_ph", all(close(rows[k]["dv1"] + rows[k]["dv2"], H["dv_tli"], 1e-12) for k in rows))
ok("  the table of dv_1, dv_2: " + "; ".join(f"k={k}: {rows[k]['dv1']:.4f} + {rows[k]['dv2']:.4f}" for k in rows),
   all(close(rows[k]["dv1"], v, 5e-5) and close(rows[k]["dv2"], w, 5e-5) for k, v, w in zip(rows, (3.0776, 3.0050, 2.9437, 2.8887, 2.8379), (0.0288, 0.1014, 0.1627, 0.2177, 0.2685))))
# the identity, symbolically: the two burns are collinear and tangential at the same perigee
v_c, v_ph, v_t = sp.symbols("v_c v_ph v_t", positive=True)
ok("dv_1 + dv_2 = (v_ph - v_c) + (v_TLI - v_ph) = v_TLI - v_c", sp.simplify((v_ph - v_c) + (v_t - v_ph) - (v_t - v_c)) == 0)
# the period condition, checked by running the clock
T3 = P3["T"]
ok(f"after 3 loops the Moon has turned {N_MOON*3*T3*R2D:.3f} deg: phi = {(phi_a + N_MOON*3*T3)*R2D:.3f} deg = phi_0", close((phi_a + N_MOON*3*T3)*R2D, PHI0*R2D, 1e-9))
# the window
cap = D - 1.5*R_SOI
Tcap = 2*np.pi*np.sqrt(((cap + R0)/2)**3/MU_E)
ok(f"apogee cap r_Moon - 1.5 r_SOI = {cap:,.0f} km: loop period at most {Tcap/DAY:.3f} d, so k loops cover waits up to {Tcap/DAY:.2f} k days",
   close(cap, 285125.6, 0.1) and close(Tcap/DAY, 6.4193, 5e-4))
kmax = {k: k*Tcap/DAY for k in (1, 2, 3, 4, 5)}
month = 2*np.pi/N_MOON/DAY
ok(f"waits up to {kmax[4]:.1f} d (k = 4) and {kmax[5]:.1f} d (k = 5) against a {month:.2f} d month: five loops cover every Moon phase", kmax[4] > 24.8 and kmax[5] > month)
ok(f"  the shortest loop that is more than the parking orbit: T = {2*np.pi*np.sqrt(R0**3/MU_E)/60:.2f} min; a wait of 2 h needs one loop of r_a = {apogee_of_period(2*HOUR):,.0f} km",
   close(2*np.pi*np.sqrt(R0**3/MU_E)/60, 90.52, 5e-3) and apogee_of_period(2*HOUR) > R0)
# lunar clearance over the loops, planar Earth + Moon, k = 3, then k = 1 (the worst)
sol3, Tk3 = propagate_loops(P3["rp"], P3["ra"], phi_a, 3, extra=4*HOUR)
dmin3, tmin3 = min_moon_distance(sol3, phi_a, t_max=3*Tk3)
ok(f"k = 3: smallest distance to the Moon during the loops {dmin3:,.0f} km = {dmin3/R_SOI:.2f} r_SOI (at t = {tmin3/DAY:.2f} d)", close(dmin3, 351724, 5) and dmin3 > 5*R_SOI)
P1 = rows[1]
sol1, Tk1 = propagate_loops(P1["rp"], P1["ra"], phi_a, 1, extra=4*HOUR)
dmin1, tmin1 = min_moon_distance(sol1, phi_a, t_max=Tk1)
ok(f"k = 1: smallest distance {dmin1:,.0f} km = {dmin1/R_SOI:.2f} r_SOI (r_a = {P1['ra']:,.0f} km)", dmin1 > 3*R_SOI)
# what the Moon does to the loops (numerical), and J2
pp = perigee_passages(sol3, Tk3, 3)
sol3_0, _ = propagate_loops(P3["rp"], P3["ra"], phi_a, 3, extra=4*HOUR, with_moon=False)
pp0 = perigee_passages(sol3_0, Tk3, 3)
ok("Earth alone: the perigee returns exactly, " + ", ".join(f"{x['r']:.3f}" for x in pp0) + " km at each pass", all(close(x["r"], P3["rp"], 1e-3) for x in pp0))
ok("Earth + Moon: perigee radius at the three returns " + ", ".join(f"{x['r']:.1f}" for x in pp) + " km (rises by " + ", ".join(f"{x['r'] - P3['rp']:.0f}" for x in pp) + " km); "
   "return times late by " + ", ".join(f"{x['dt']:.0f}" for x in pp) + " s",
   close(pp[0]["r"] - P3["rp"], 41.4, 0.5) and close(pp[1]["r"] - P3["rp"], 70.9, 0.5) and close(pp[2]["r"] - P3["rp"], 51.5, 0.5) and close(pp[0]["dt"], 77.8, 2.0))
ok("  the perigee direction at the three returns (the third is the trans-lunar burn) moves by " + ", ".join(f"{x['angle']*R2D:.4f}" for x in pp) + " deg (Moon alone)",
   len(pp) == 3 and close(pp[0]["angle"]*R2D, 0.0076, 5e-4) and close(pp[1]["angle"]*R2D, -0.0109, 5e-4) and close(pp[2]["angle"]*R2D, -0.0395, 5e-4))
pla = P3["a"]*(1 - P3["e"]**2)
inc = 28.5*D2R
wd = j2_apse_rate(pla, P3["n"], inc)
Od = L.j2_node_rate(pla, P3["n"], inc)
ok(f"J2 at i = 28.5 deg, p = {pla:,.0f} km: the apse turns {wd*R2D*DAY:.4f} deg/d, the node regresses {Od*R2D*DAY:.4f} deg/d; over the {W/DAY:.2f} d wait {wd*W*R2D:.3f} and {Od*W*R2D:.3f} deg",
   close(wd*R2D*DAY, 0.0943, 5e-4) and close(Od*R2D*DAY, -0.0579, 5e-4) and close(wd*W*R2D, 0.629, 2e-3) and close(Od*W*R2D, -0.387, 2e-3))


def pvec(om, Om, i_):                                   # perigee direction in space
    return np.array([np.cos(om)*np.cos(Om) - np.sin(om)*np.cos(i_)*np.sin(Om), np.cos(om)*np.sin(Om) + np.sin(om)*np.cos(i_)*np.cos(Om), np.sin(om)*np.sin(i_)])


nvec = lambda Om, i_: np.array([np.sin(i_)*np.sin(Om), -np.sin(i_)*np.cos(Om), np.cos(i_)])
rots = np.array([np.degrees(np.arccos(np.clip(pvec(o*D2R, 0.0, inc) @ pvec(o*D2R + wd*W, Od*W, inc), -1, 1))) for o in np.arange(0, 360, 15)])
ok(f"  the perigee direction turns in space by {rots.min():.3f} to {rots.max():.3f} deg over the wait (depending on where the perigee is in the plane), i.e. {rots.min()*D2R/N_MOON/HOUR:.2f} to {rots.max()*D2R/N_MOON/HOUR:.2f} h of the Moon's motion",
   close(rots.min(), 0.290, 2e-3) and close(rots.max(), 0.343, 2e-3) and close(rots.min()*D2R/N_MOON/HOUR, 0.527, 5e-3) and close(rots.max()*D2R/N_MOON/HOUR, 0.625, 5e-3))
tilt = np.degrees(np.arccos(np.clip(nvec(0.0, inc) @ nvec(Od*W, inc), -1, 1)))
ok(f"  the plane tilts by {tilt:.3f} deg (|sin i dOmega| = {abs(np.degrees(np.sin(inc)*Od*W)):.3f}): {D*np.sin(tilt*D2R):,.0f} km out of plane at the Moon's distance",
   close(tilt, 0.184, 2e-3) and close(D*np.sin(tilt*D2R), 1237, 10))
# phasing correction by a small perigee burn (cheapest place, the manoeuvre notes, sec-phasing)
dT = 600.0
aph2 = P3["a"]*(1 + dT/P3["T"])**(2/3)
dvcorr = vis_viva(R0, aph2) - vis_viva(R0, P3["a"])
ok(f"a perigee burn of {1000*dvcorr:.1f} m/s lengthens the loop by 10 min (the loop is then {P3['T']/DAY*24:.2f} -> {(P3['T'] + dT)/DAY*24:.2f} h)", close(1000*dvcorr, 0.538, 0.002))

# =========================================================================================== Part 3
print("Part 3: the low-thrust spiral")
# -- the derivations
r, mu_, f, t_ = sp.symbols("r mu f t", positive=True)
v_t_ = sp.Function("v")(t_)
eps_c = -mu_/(2*r)
# circular orbit: v^2 = mu/r; d(eps)/dt = f v with eps = -v^2/2  ->  dv/dt = -f
vdot_ = sp.symbols("vdot")
ok("circular orbit: eps = -v^2/2 so d eps/dt = -v vdot; setting it equal to f v gives vdot = -f", sp.solve(sp.Eq(-v_t_*vdot_, f*v_t_), vdot_) == [-f])
v1_, v2_, vv = sp.symbols("v1 v2 vv", positive=True)
Nrev = sp.integrate(vv**3/(2*sp.pi*mu_*f), (vv, v2_, v1_))
ok("revolutions N = (v1^4 - v2^4)/(8 pi mu f)", sp.simplify(Nrev - (v1_**4 - v2_**4)/(8*sp.pi*mu_*f)) == 0)
eps1 = sp.symbols("eps1", positive=True)
r1s, r2s = sp.symbols("r1 r2", positive=True)
Nr = (v1_**4 - v2_**4)/(8*sp.pi*mu_*f)
Nr_r = Nr.subs({v1_: sp.sqrt(mu_/r1s), v2_: sp.sqrt(mu_/r2s)})
ok("  = (1/(8 pi eps_1)) (1 - (r_1/r_2)^2) with eps_1 = a r_1^2/mu", sp.simplify(Nr_r - (1 - (r1s/r2s)**2)*mu_/(8*sp.pi*f*r1s**2)) == 0)
# the eccentricity vector equation, for any perturbing acceleration (planar)
x, y, u, w, ax, ay = sp.symbols("x y u w a_x a_y", real=True)
rr = sp.sqrt(x**2 + y**2)
h = x*w - y*u
evec = sp.Matrix([w*h, -u*h])/mu_ - sp.Matrix([x, y])/rr
state = [x, y, u, w]
deriv = [u, w, -mu_*x/rr**3 + ax, -mu_*y/rr**3 + ay]
dev = sp.Matrix([sum(sp.diff(evec[i], s)*d for s, d in zip(state, deriv)) for i in range(2)])
rv, vv_, av = sp.Matrix([x, y]), sp.Matrix([u, w]), sp.Matrix([ax, ay])
claim = (2*(vv_.dot(av))*rv - (rv.dot(av))*vv_ - (rv.dot(vv_))*av)/mu_
ok("d e/dt = [2 (v.a_d) r - (r.a_d) v - (r.v) a_d]/mu  (any perturbing acceleration)", sp.simplify(dev - claim) == sp.zeros(2, 1))
# circular orbit, tangential thrust: d e/dt = (2 a_T / v) r_hat, a circle of radius 2 eps traced at the mean motion
nn_, aT_, tt = sp.symbols("n a_T tt", positive=True)
v_circ = sp.sqrt(mu_/r)
ev_tangential = claim.subs({x: r, y: 0, u: 0, w: v_circ, ax: 0, ay: aT_})
ok("circular orbit, tangential a_T:  d e/dt = (2 a_T/v) r_hat", sp.simplify(ev_tangential - sp.Matrix([2*aT_/v_circ, 0])) == sp.zeros(2, 1))
# integrate (2 a_T/v)(cos nt, sin nt) with n = v/r:  e = (2 a_T/(v n)) (sin nt, 1 - cos nt); |e| = 4 eps |sin(nt/2)|, eps = a_T r^2/mu
n_c = sp.sqrt(mu_/r**3)
amp = sp.simplify(2*aT_/(v_circ*n_c))
ok("amplitude 2 a_T/(v n) = 2 eps with eps = a_T r^2/mu", sp.simplify(amp - 2*aT_*r**2/mu_) == 0)
phi_ = sp.symbols("phi", real=True)
ok("|e| = 2 eps sqrt(sin^2 + (1 - cos)^2) = 4 eps |sin(nt/2)|", sp.simplify(sp.sqrt(sp.sin(phi_)**2 + (1 - sp.cos(phi_))**2) - 2*sp.Abs(sp.sin(phi_/2))) == 0
   or all(close(float(sp.sqrt(sp.sin(q)**2 + (1 - sp.cos(q))**2)), 2*abs(np.sin(q/2)), 1e-12) for q in np.linspace(0, 6.28, 25)))

# -- numerical integration at a SMART-1-class acceleration
a_avg = 2.737/(3648*3600)                           # km/s^2, from ESA's 26 Oct 2004 status report: 2737 m/s in 3648 h of thrusting
ok(f"SMART-1 (ESA status report of 26 Oct 2004: 2737 m/s in 3648 h of firing): average thrust acceleration {a_avg*1e3:.4e} m/s^2 = {a_avg:.4e} km/s^2",
   close(a_avg, 2.0841e-7, 5e-11))
aT = 2.08e-7
r_end = 1.0e5
sol = integrate_spiral(R0, aT, r_end)
t_num = sol.t[-1]
t_an = spiral_time(R0, r_end, aT)
ok(f"spiral from 6678 km to 100 000 km at a_T = 2.08e-7 km/s^2: dv = {spiral_dv(R0, r_end):.4f} km/s, analytic time {t_an/DAY:.3f} d, integrated {t_num/DAY:.3f} d ({100*abs(t_num - t_an)/t_an:.4f}% apart)",
   close(spiral_dv(R0, r_end), 5.7293, 5e-5) and close(t_an/DAY, 318.802, 5e-3) and abs(t_num - t_an)/t_an < 2e-4)
ts = np.linspace(0, t_num, 120001)
Y = sol.sol(ts)
th = np.unwrap(np.arctan2(Y[1], Y[0]))
rev_num = (th[-1] - th[0])/(2*np.pi)
rev_an = spiral_revolutions(R0, r_end, aT)
eps_end = aT*r_end**2/MU_E
ok(f"revolutions: integrated {rev_num:.2f}, formula (v_1^4 - v_2^4)/(8 pi mu a_T) = {rev_an:.2f}", close(rev_num, rev_an, 0.02) and close(rev_an, 1702.09, 0.01))
eps_loc = aT*np.hypot(Y[0], Y[1])**2/MU_E
idx = np.arange(0, len(ts), 40)
e_os = np.array([osculating_e(Y[:, i]) for i in idx])
ratio = e_os/(4*eps_loc[idx])
dev = np.abs(e_os - 2*eps_loc[idx])
eps1_ = aT*R0**2/MU_E
ok(f"osculating e: never above 4 kappa(r) (max e/(4 kappa) = {ratio.max():.3f}); e is within 2 kappa_1 of 2 kappa(r) throughout: max |e - 2 kappa|/(2 kappa_1) = {dev.max()/(2*eps1_):.4f}; at the end e = {e_os[-1]:.5f}, 2 kappa = {2*eps_end:.5f}",
   ratio.max() < 1.0 and dev.max() <= 2*eps1_*(1 + 1e-3) and abs(e_os[-1] - 2*eps_end) < 2*eps1_ and close(eps_loc[0], 2.3272e-5, 5e-9))
rr_all = np.hypot(Y[0], Y[1]); vv_all = np.hypot(Y[2], Y[3])
vr_all = (Y[0]*Y[2] + Y[1]*Y[3])/rr_all
dvr = np.abs(vr_all - 2*(aT*rr_all**2/MU_E)*vv_all)[idx]
ok(f"the steady eccentricity is the climb rate: radial speed v_r = 2 kappa v (flight-path angle 2 kappa) to within 2 kappa_1 v: max |v_r - 2 kappa v|/(2 kappa_1 v) = {np.max(dvr/(2*eps1_*vv_all[idx])):.4f}",
   np.max(dvr/(2*eps1_*vv_all[idx])) <= 1 + 1e-3)
rs_, mu_s, aT_s = sp.symbols("r mu a_T", positive=True)
ok("2 kappa v with kappa = a_T r^2/mu and v = sqrt(mu/r) is 2 a_T r^(3/2)/sqrt(mu), the rdot of the energy equation",
   sp.simplify(2*(aT_s*rs_**2/mu_s)*sp.sqrt(mu_s/rs_) - 2*aT_s*rs_**(sp.Rational(3, 2))/sp.sqrt(mu_s)) == 0)
ok(f"the speed is below the circular speed, as thrust forward slows the craft: at the end v = {np.hypot(*Y[2:, -1]):.4f}, v_c(r) = {np.sqrt(MU_E/np.hypot(*Y[:2, -1])):.4f} km/s",
   close(np.hypot(*Y[2:, -1]), 1.99672, 5e-5))
# radius against time: analytic v_c(t) = v_c1 - a_T t
rt = spiral_r_of_t(ts, R0, aT)
rn = np.hypot(Y[0], Y[1])
ok(f"r(t) = mu/(v_c1 - a_T t)^2 agrees with the integration: largest relative difference {np.max(np.abs(rt - rn)/rn):.2e}", np.max(np.abs(rt - rn)/rn) < 3e-3)
# one revolution at constant eps: the 4 eps bound is attained
solc = integrate_spiral(R0, 1e-4*MU_E/R0**2, 1.01*R0)   # eps_1 = 1e-4; only the first 2 revolutions are used
tc = np.linspace(0, 2*2*np.pi*np.sqrt(R0**3/MU_E), 4001)
Yc = solc.sol(tc)
ec = np.array([osculating_e(Yc[:, i]) for i in range(len(tc))])
ok(f"first two revolutions at kappa_1 = 1e-4: e oscillates between {ec.min():.2e} and {ec.max():.2e} (4 kappa = {4e-4:.2e})", close(ec.max(), 4e-4, 2e-5))


def evec(y):
    rv_, vv_ = np.array(y[:2]), np.array(y[2:])
    hh = rv_[0]*vv_[1] - rv_[1]*vv_[0]
    return np.array([vv_[1]*hh, -vv_[0]*hh])/MU_E - rv_/np.hypot(*rv_)


th_c = np.unwrap(np.arctan2(Yc[1], Yc[0]))
evc = np.array([evec(Yc[:, i]) for i in range(len(tc))])
pred = 2e-4*np.stack([np.sin(th_c), 1 - np.cos(th_c)], axis=1)
ok(f"the vector, not only its length: e = 2 kappa_1 (sin theta, 1 - cos theta) for constant kappa; largest component difference {np.max(np.abs(evc - pred)):.1e} (kappa grows by 1.3% in two revolutions: 2 kappa x 0.013 = 2.6e-6)",
   np.max(np.abs(evc - pred)) < 3e-6)
bypart = lambda th_, k_, k1_: 2*k_[:, None]*np.stack([np.sin(th_), -np.cos(th_)], axis=1) + 2*k1_*np.array([0.0, 1.0])
kc = 1e-4*(np.hypot(Yc[0], Yc[1])/R0)**2
ok(f"with kappa(r) growing: e = 2 kappa(r) (sin theta, -cos theta) + 2 kappa_1 j + O(kappa^2); largest component difference {np.max(np.abs(evc - bypart(th_c, kc, 1e-4))):.1e} (2 kappa = 2e-4: 0.13%)",
   np.max(np.abs(evc - bypart(th_c, kc, 1e-4))) < 3e-7 and np.max(np.abs(evc - bypart(th_c, kc, 1e-4))) < 0.25*np.max(np.abs(evc - pred)))
# -- where the model stops
eps_lim = 0.01
r_lim = np.sqrt(eps_lim*MU_E/aT)
tide_ratio = lambda rr_: tidal_ratio(rr_)/aT
ok(f"eps = 0.01 at r = {r_lim:,.0f} km (a_T = 2.08e-7 km/s^2); eps at 100 000 km = {eps_end:.4f}, at the Moon's distance = {aT*D**2/MU_E:.3f}",
   close(r_lim, 138410, 50) and close(aT*D**2/MU_E, 0.0771, 5e-4))
ok(f"at the Moon's distance kappa = {aT*D**2/MU_E:.4f}: the flight-path angle 2 kappa = {2*aT*D**2/MU_E:.3f} rad = {np.degrees(2*aT*D**2/MU_E):.1f} deg", close(2*aT*D**2/MU_E, 0.154, 5e-4) and close(np.degrees(2*aT*D**2/MU_E), 8.8, 0.05))
ok(f"the Moon's tidal acceleration over the thrust acceleration: {tide_ratio(1e5):.3f} at 100 000 km, {tide_ratio(2e5):.3f} at 200 000 km; equal to a_T only at r ~ {brentq(lambda rr_: tidal_ratio(rr_) - aT, 1e5, 3.4e5):,.0f} km",
   close(tide_ratio(1e5), 0.1319, 5e-4) and close(tide_ratio(2e5), 0.5337, 5e-4) and close(brentq(lambda rr_: tidal_ratio(rr_) - aT, 1e5, 3.4e5), 241822, 5))
d_pull = np.sqrt(MU_M/aT)
ok(f"the Moon's own pull equals the thrust acceleration within {d_pull:,.0f} km of its centre = {d_pull/R_SOI:.2f} r_SOI", close(d_pull, 153537, 20))
dv_at = lambda r_: spiral_dv(R0, r_)
ok(f"share of the whole spiral to the Moon's distance that the Earth-only model covers up to 100 000 km: dv {100*dv_at(1e5)/dv_at(D):.1f}% (of {dv_at(D):.4f} km/s), time the same share",
   close(dv_at(D), 6.7075, 5e-5) and close(100*dv_at(1e5)/dv_at(D), 85.4, 0.05))
# -- rocket and power
c_ = G0*1800
ok(f"c = g0 Isp = {c_:.3f} km/s at Isp = 1800 s; propellant fraction for dv = 5.7293 km/s: {100*propellant_fraction(5.7293, 1800):.1f}%", close(c_, 17.652, 5e-4) and close(100*propellant_fraction(5.7293, 1800), 27.7, 0.05))
P_w, eta_ = 5000.0, 0.6
F18 = thrust_from_power(P_w, 1800, eta_)
ok(f"5 kW, efficiency 0.6, Isp 1800 s: F = 2 eta P/(g0 Isp) = {F18:.4f} N; initial acceleration {F18/1000*1e-3:.3e} km/s^2 at 1000 kg", close(F18, 0.33992, 5e-5))
t_ct = trip_time_constant_thrust(1000.0, F18, 1800, 5.7293)
t_ca = 5.7293/(F18/1000*1e-3)
ok(f"time with constant thrust {t_ct/DAY:.2f} d = (c m_0/F)(1 - exp(-dv/c)); with the acceleration frozen at its initial value {t_ca/DAY:.2f} d", close(t_ct/DAY, 166.60, 0.05) and close(t_ca/DAY, 195.0, 0.5))
ok(f"  propellant {1000*propellant_fraction(5.7293, 1800):.0f} kg of 1000 kg; final mass {1000 - 1000*propellant_fraction(5.7293, 1800):.0f} kg",
   close(1000*propellant_fraction(5.7293, 1800), 277.2, 0.1))
# the Isp trade at fixed power (no power-system mass)
rows_isp = []
for isp in (500, 1000, 1800, 3000, 6000):
    Fi = thrust_from_power(P_w, isp, eta_)
    rows_isp.append((isp, Fi, 100*propellant_fraction(5.7293, isp), trip_time_constant_thrust(1000.0, Fi, isp, 5.7293)/DAY))
ok("fixed power: " + "; ".join(f"{i}s: F={F_:.3f} N, {pf:.1f}%, {td:.0f} d" for i, F_, pf, td in rows_isp),
   all(rows_isp[i][3] < rows_isp[i + 1][3] for i in range(4)) and all(rows_isp[i][2] > rows_isp[i + 1][2] for i in range(4)))

# =========================================================================================== Part 4
print("Part 4: worked examples and comparisons")
# like-for-like at the end of the model: a circle of radius 100 000 km
aH = 0.5*(R0 + 1e5)
dvH1 = vis_viva(R0, aH) - V_C0
dvH2 = np.sqrt(MU_E/1e5) - vis_viva(1e5, aH)
tH = np.pi*np.sqrt(aH**3/MU_E)
ok(f"Hohmann LEO -> 100 000 km circle: dv {dvH1:.4f} + {dvH2:.4f} = {dvH1 + dvH2:.4f} km/s in {tH/HOUR:.2f} h; propellant at Isp 320 s {100*propellant_fraction(dvH1 + dvH2, 320):.1f}%",
   close(dvH1 + dvH2, 4.1425, 5e-4) and close(tH/HOUR, 17.03, 5e-3) and close(100*propellant_fraction(dvH1 + dvH2, 320), 73.3, 0.1))
ok(f"  the spiral costs {100*(dv_at(1e5)/(dvH1 + dvH2) - 1):.0f}% more dv ({dv_at(1e5):.4f} vs {dvH1 + dvH2:.4f} km/s) and {t_ct/tH:.0f} times the time", close(100*(dv_at(1e5)/(dvH1 + dvH2) - 1), 38.3, 0.1))
# the lunar mission, like for like: capture into a 100 km circular lunar orbit
rl = RM + 100.0
v_low = np.sqrt(MU_M/rl)
dv_from_soi = v_low - np.sqrt(MU_M/R_SOI)
r_valid_moon = np.sqrt(eps_lim*MU_M/aT)
dv_from_valid = v_low - np.sqrt(MU_M/r_valid_moon)
ok(f"spiral down about the Moon: from the SOI edge dv = {dv_from_soi:.4f} km/s, from r = {r_valid_moon:,.0f} km (eps = 0.01 there) dv = {dv_from_valid:.4f} km/s; the circular speed at 100 km is {v_low:.4f}",
   close(dv_from_soi, 1.3613, 5e-4) and close(dv_from_valid, 1.0685, 5e-4) and close(r_valid_moon, 15357, 10) and close(v_low, 1.6336, 5e-4))
ok(f"  ideal all-electric: {dv_at(D):.4f} + {dv_from_soi:.4f} = {dv_at(D) + dv_from_soi:.4f} km/s; propellant at Isp 1800 s {100*propellant_fraction(dv_at(D) + dv_from_soi, 1800):.1f}%; time at 2.08e-7 km/s^2 {(dv_at(D) + dv_from_soi)/aT/DAY:.0f} d",
   close(dv_at(D) + dv_from_soi, 8.0688, 5e-4) and close(100*propellant_fraction(dv_at(D) + dv_from_soi, 1800), 36.7, 0.1) and close((dv_at(D) + dv_from_soi)/aT/DAY, 449.1, 0.5))
dv_direct = H["dv_tli"] + 0.7949
ok(f"  impulsive (Example 2a) {dv_direct:.4f} km/s; propellant at Isp 320 s {100*propellant_fraction(dv_direct, 320):.1f}%; loop adds only time: {W/DAY:.2f} d waiting",
   close(dv_direct, 3.9013, 5e-5) and close(100*propellant_fraction(dv_direct, 320), 71.1, 0.1))
ok(f"  Hohmann to the Moon's orbit and a circle there: {H['dv_tli']:.4f} + {np.sqrt(MU_E/D) - vis_viva(D, H['a']):.4f} = {H['dv_tli'] + np.sqrt(MU_E/D) - vis_viva(D, H['a']):.4f} km/s (the spiral to the same circle costs {dv_at(D):.4f}, {100*(dv_at(D)/(H['dv_tli'] + np.sqrt(MU_E/D) - vis_viva(D, H['a'])) - 1):.0f}% more)",
   close(H["dv_tli"] + np.sqrt(MU_E/D) - vis_viva(D, H["a"]), 3.9365, 5e-5) and close(100*(dv_at(D)/(H["dv_tli"] + np.sqrt(MU_E/D) - vis_viva(D, H["a"])) - 1), 70.4, 0.1))
# the number of revolutions and the exaggerated spiral of the figure
ok(f"figure with eps_1 = 0.003: {spiral_revolutions(R0, 6*R0, 0.003*MU_E/R0**2):.2f} revolutions from r_1 to 6 r_1", close(spiral_revolutions(R0, 6*R0, 0.003*MU_E/R0**2), 12.9, 0.05))
# Edelbaum for the record (derived in complications-to-impulsive-transfers-chobotov.qmd)
V1_, V2_ = np.sqrt(MU_E/R0), np.sqrt(MU_E/1e5)
edel = lambda al: np.sqrt(V1_**2 + V2_**2 - 2*V1_*V2_*np.cos(np.pi/2*al*D2R))
ok(f"Edelbaum's plane-change term (derived in complications-to-impulsive-transfers-chobotov.qmd): LEO to r = 100 000 km with 10 deg of plane change {edel(10):.3f} km/s, "
   f"28.5 deg {edel(28.5):.3f} km/s, against {dv_at(1e5):.3f} km/s with none",
   close(edel(10), 5.829, 5e-4) and close(edel(28.5), 6.463, 5e-4) and close(edel(0), dv_at(1e5), 1e-12))
# -- the energy bought per unit of dv: why the spiral costs more than the Hohmann transfer
dE = 0.5*MU_E*(1/R0 - 1/1e5)
vb1, vb2 = 0.5*(V_C0 + vis_viva(R0, aH)), 0.5*(vis_viva(1e5, aH) + np.sqrt(MU_E/1e5))
ok(f"an impulse buys energy at the mean of its before and after speeds (d eps = v dv + dv^2/2): burn 1 at {vb1:.3f} km/s ({V_C0:.3f} -> {vis_viva(R0, aH):.3f}), burn 2 at {vb2:.3f} ({vis_viva(1e5, aH):.3f} -> {np.sqrt(MU_E/1e5):.3f}); "
   f"their energies {vb1*dvH1:.3f} + {vb2*dvH2:.3f} = {vb1*dvH1 + vb2*dvH2:.3f}",
   close(vb1, 9.152, 5e-4) and close(vb2, 1.3515, 5e-4) and close(vb1*dvH1 + vb2*dvH2, 0.5*MU_E*(1/R0 - 1/1e5), 1e-9))
ok(f"energy to be gained {dE:.3f} km^2/s^2 = integral of v d(dv): mean speed at which it is bought {dE/dv_at(1e5):.3f} km/s in the spiral, {dE/(dvH1 + dvH2):.3f} km/s in the Hohmann transfer",
   close(dE, 27.851, 5e-4) and close(dE/dv_at(1e5), 4.861, 5e-4) and close(dE/(dvH1 + dvH2), 6.723, 5e-4))
# -- the task and the exercises
phi0_ = PHI0
Wm = required_wait(phi0_, -60*D2R)
Pm = {k: phasing_orbit(Wm/k) for k in (1, 2, 3, 4)}
ok(f"task (Moon 60 deg behind): W = {Wm/DAY:.3f} d; r_a for k = 1..4: " + ", ".join(f"{Pm[k]['ra']:,.0f}" for k in Pm) + f" km; cap {cap:,.0f} km: k >= 3",
   close(Wm/DAY, 13.875, 5e-4) and all(close(Pm[k]["ra"], v, 1.5) for k, v in zip(Pm, (481131, 300622, 227836, 186909)))
   and Pm[2]["ra"] > cap and Pm[3]["ra"] < cap and close(Pm[2]["T"]/DAY, 6.937, 5e-4) and close(Pm[3]["T"]/DAY, 4.625, 5e-4) and close(Pm[4]["T"]/DAY, 3.469, 5e-4))
W1 = required_wait(phi0_, 60*D2R)
P1x = {k: phasing_orbit(W1/k) for k in (1, 2, 3)}
ok(f"exercise 1 (Moon 60 deg ahead): W = {W1/DAY:.3f} d; k = 1, 2, 3: T = " + ", ".join(f"{P1x[k]['T']/DAY:.3f}" for k in P1x) + " d, r_a = " + ", ".join(f"{P1x[k]['ra']:,.0f}" for k in P1x) + " km",
   close(W1/DAY, 4.780, 5e-4) and all(close(P1x[k]["ra"], v, 1.5) for k, v in zip(P1x, (233047, 144339, 108570))) and all(P1x[k]["ra"] < cap for k in P1x))
ok(f"exercise 3: 4 eps = 0.1 at r = {np.sqrt(0.025*MU_E/aT):,.0f} km", close(np.sqrt(0.025*MU_E/aT), 218881, 1))
F10 = thrust_from_power(10000.0, 1800, 0.6)
ok(f"exercise 5: 10 kW: F = {F10:.3f} N, same propellant, t = {trip_time_constant_thrust(1000.0, F10, 1800, 5.7293)/DAY:.1f} d", close(F10, 0.680, 5e-4) and close(trip_time_constant_thrust(1000.0, F10, 1800, 5.7293)/DAY, 83.3, 0.05))
t_down = dv_from_valid_15350 = (v_low - np.sqrt(MU_M/15350.0))/aT/DAY
ok(f"exercise 6: spiral down from 15 350 km to 100 km altitude: {t_down:.1f} d", close(t_down, 59.4, 0.05))
F30 = thrust_from_power(5000.0, 3000, 0.6)
ok(f"task of Example 2 at 3000 s: F = {F30:.3f} N, propellant {100*propellant_fraction(5.7293, 3000):.1f}%, t = {trip_time_constant_thrust(1000.0, F30, 3000, 5.7293)/DAY:.0f} d; 10 kg less per 100 kg and {trip_time_constant_thrust(1000.0, F30, 3000, 5.7293)/DAY - t_ct/DAY:.0f} d more",
   close(F30, 0.204, 5e-4) and close(100*propellant_fraction(5.7293, 3000), 17.7, 0.05) and close(trip_time_constant_thrust(1000.0, F30, 3000, 5.7293)/DAY, 295, 1)
   and close(trip_time_constant_thrust(1000.0, F30, 3000, 5.7293)/DAY - t_ct/DAY, 128.4, 0.5))
ok(f"mass of propellant: spiral {1000*propellant_fraction(5.7293, 1800):.0f} kg, chemical {1000*propellant_fraction(dvH1 + dvH2, 320):.0f} kg ({propellant_fraction(dvH1 + dvH2, 320)/propellant_fraction(5.7293, 1800):.1f} times more); c ratio {1800/320:.1f}",
   close(propellant_fraction(dvH1 + dvH2, 320)/propellant_fraction(5.7293, 1800), 2.64, 0.01))
ok(f"minutes of the Moon's motion in one parking orbit: {N_MOON*R2D*(2*np.pi*np.sqrt(R0**3/MU_E))*1.0:.2f} deg per 90.5 min", close(N_MOON*R2D*2*np.pi*np.sqrt(R0**3/MU_E), 0.83, 0.01))
print("done")