"""Checks the numbers quoted in cr3bp.qmd, using the helper module _cr3bp.py.

Run from the sfm-slides directory:   python3 codes/cr3bp-verify.py
"""
import sys
import warnings
import numpy as np
from scipy.optimize import brentq
sys.path.insert(0, "codes")
from _cr3bp import (SYSTEMS, units, omega, omega_grad, rhs, jacobi, collinear, triangular, hessian, linear_eigenvalues,
                    c2_of, collinear_linear, propagate, lyapunov, halo, orbit_history, _next_crossing, DAY)

warnings.simplefilter("ignore")
RESULTS = []


def check(name, value, stated, tol):
    ok = abs(value - stated) <= tol
    RESULTS.append(ok)
    print(f"{'PASS' if ok else 'FAIL'}  {name:<66} {value:16.8g}   page: {stated}")


def truth(name, cond):
    RESULTS.append(bool(cond))
    print(f"{'PASS' if cond else 'FAIL'}  {name}")


mu_em, L_em, T_em, V_em = units("Earth-Moon")
mu_se, L_se, T_se, V_se = units("Sun-(Earth+Moon)")

print("Units")
check("Earth-Moon mu", mu_em, 0.0121506, 5e-8)
check("  time unit [d]", T_em / DAY, 4.3425, 5e-5)
check("  time unit [s]", T_em, 375190, 1)
check("  speed unit [km/s]", V_em, 1.0245, 5e-5)
check("Sun-Earth mu", mu_se, 3.0404e-6, 5e-11)
check("  time unit [d]", T_se / DAY, 58.132, 5e-4)
check("  speed unit [km/s]", V_se, 29.785, 5e-4)

print("Equations of motion and the Jacobi constant")
rng = np.random.default_rng(1)
s0 = np.array([0.3, 0.4, 0.1, 0.2, -0.3, 0.15])
h = 1e-6
grad_fd = [(omega(*(s0[:3] + h * e), mu_em) - omega(*(s0[:3] - h * e), mu_em)) / (2 * h) for e in np.eye(3)]
check("gradient of Omega against finite differences", np.abs(np.array(omega_grad(s0, mu_em)) - np.array(grad_fd)).max(), 0, 1e-8)
f = rhs(0, s0, mu_em)
check("x equation:  xdd - 2 ydot = Omega_x", f[3] - 2 * s0[4], omega_grad(s0, mu_em)[0], 1e-14)
check("y equation:  ydd + 2 xdot = Omega_y", f[4] + 2 * s0[3], omega_grad(s0, mu_em)[1], 1e-14)
# C = 2 (h_z - E) with inertial energy and angular momentum
x, y, z, vx, vy, vz = s0
vI = np.array([vx - y, vy + x, vz])
r1 = np.sqrt((x + mu_em) ** 2 + y * y + z * z); r2 = np.sqrt((x - 1 + mu_em) ** 2 + y * y + z * z)
E_in = 0.5 * vI @ vI - (1 - mu_em) / r1 - mu_em / r2
hz = x * vI[1] - y * vI[0]
check("C = 2 (h_z - E_inertial)", jacobi(s0, mu_em), 2 * (hz - E_in), 1e-12)
ts, Y = orbit_history(mu_em, [0.8, 0.1, 0.05, 0.1, 0.2, 0.0], 12.0, 3000)
Cs = np.array([jacobi(Y[:, k], mu_em) for k in range(Y.shape[1])])
check("Jacobi constant conserved along an integration (spread)", np.ptp(Cs), 0, 1e-10)

print("Equilibrium points")
for name, (mu, L, T, V) in {"Earth-Moon": (mu_em, L_em, T_em, V_em), "Sun-Earth": (mu_se, L_se, T_se, V_se)}.items():
    L1, L2, L3 = collinear(mu)
    for lab, xp in (("L1", L1), ("L2", L2), ("L3", L3)):
        check(f"{name} {lab}: Omega_x = 0", omega_grad([xp, 0, 0, 0, 0, 0], mu)[0], 0, 1e-12)
    L4, L5 = triangular(mu)
    check(f"{name} L4: Omega_x", omega_grad(list(L4) + [0, 0, 0], mu)[0], 0, 1e-12)
    check(f"{name} L4: Omega_y", omega_grad(list(L4) + [0, 0, 0], mu)[1], 0, 1e-12)
    check(f"{name} L4: r1 = r2 = 1", np.hypot(L4[0] + mu, L4[1]), 1.0, 1e-12)
    check(f"{name} L4: r2", np.hypot(L4[0] - 1 + mu, L4[1]), 1.0, 1e-12)
check("Earth-Moon L1 x", collinear(mu_em)[0], 0.836915, 5e-7)
check("Earth-Moon L2 x", collinear(mu_em)[1], 1.155682, 5e-7)
check("Earth-Moon L3 x", collinear(mu_em)[2], -1.005063, 5e-7)
check("Sun-Earth L1 x", collinear(mu_se)[0], 0.989986, 5e-7)
check("Sun-Earth L2 x", collinear(mu_se)[1], 1.010075, 5e-7)
check("Sun-Earth L3 x", collinear(mu_se)[2], -1.000001, 5e-7)
check("L3 approximation -(1 + 5 mu/12)", -(1 + 5 * mu_em / 12), collinear(mu_em)[2], 5e-6)
g1 = 1 - mu_em - collinear(mu_em)[0]; g2 = collinear(mu_em)[1] - (1 - mu_em)
check("Earth-Moon gamma_1 [units]", g1, 0.150934, 5e-7); check("  gamma_2", g2, 0.167833, 5e-7)
check("  gamma_1 [km]", g1 * L_em, 58019, 1); check("  gamma_2 [km]", g2 * L_em, 64515, 1)
check("  L1 from the Earth [km]", (collinear(mu_em)[0] + mu_em) * L_em, 326381, 1)
check("Hill estimate (mu/3)^(1/3), Earth-Moon", (mu_em / 3) ** (1 / 3), 0.159401, 5e-7)
check("  [km]", (mu_em / 3) ** (1 / 3) * L_em, 61274, 1)
check("  mean of gamma_1, gamma_2", (g1 + g2) / 2, 0.159384, 5e-7)
check("Hill estimate, Sun-Earth [AU]", (mu_se / 3) ** (1 / 3), 0.0100447, 5e-8)
check("  [km]", (mu_se / 3) ** (1 / 3) * L_se, 1.5027e6, 500)
check("Sun-Earth L1 from the Earth [km]", (1 - mu_se - collinear(mu_se)[0]) * L_se, 1497621, 1)
check("Sun-Earth L2 from the Earth [km]", (collinear(mu_se)[1] - 1 + mu_se) * L_se, 1507683, 1)
for name, mu, st in (("Earth-Moon", mu_em, (3.188341, 3.172160, 3.012147, 2.987997)), ("Sun-Earth", mu_se, (3.000898, 3.000894, 3.000003, 2.999997))):
    Lc = collinear(mu); L4 = triangular(mu)[0]
    Cs_ = [jacobi([p, 0, 0, 0, 0, 0], mu) for p in Lc] + [jacobi([L4[0], L4[1], 0, 0, 0, 0], mu)]
    for lab, v_, s_ in zip(("C_L1", "C_L2", "C_L3", "C_L4"), Cs_, st):
        check(f"{name} {lab}", v_, s_, 5e-7)
    check(f"{name} C_L4 = 3 - mu(1 - mu)", Cs_[3], 3 - mu * (1 - mu), 1e-12)
# quintic degree: the numerator of f has degree 5
import sympy as sp
xs, ms = sp.symbols("x mu", real=True)
for lo, hi, lab in ((-1, 1, "L1 interval"),):
    pass
fx = xs - (1 - ms) / (xs + ms) ** 2 - ms / (xs - 1 + ms) ** 2           # for x > 1 - mu (both distances positive)
num = sp.numer(sp.together(fx))
truth("the collinear equation is a quintic in x (right of both primaries)", sp.degree(sp.expand(num), xs) == 5)

print("Time-of-flight task (opening the L1 neck)")
n_em = 1 / T_em
rp = 6378.137 + 300
x_p = -mu_em + rp / L_em
Cl1 = jacobi([collinear(mu_em)[0], 0, 0, 0, 0, 0], mu_em)
Cfun = lambda vin: jacobi([x_p, 0, 0, 0, (vin - n_em * rp) / V_em, 0], mu_em)
v_open = brentq(lambda v: Cfun(v) - Cl1, 7.7, 13)
check("perigee speed that opens the L1 neck [km/s]", v_open, 10.79, 0.005)
vc = np.sqrt(398600.4418 / rp)
check("  escape speed [km/s]", np.sqrt(2) * vc, 10.926, 5e-4)
a_h = (6678.137 + 384400) / 2
check("  Hohmann perigee speed to the Moon's distance [km/s]", np.sqrt(398600.4418 * (2 / 6678.137 - 1 / a_h)), 10.83, 0.005)
check("  rotating-frame correction n r_p [km/s]", n_em * rp, 0.0178, 5e-5)
check("  n [rad/s]", n_em, 2.6653e-6, 5e-11)

print("Stability")
L1e, L2e, L3e = collinear(mu_em)
for name, mu, xp, st in (("Earth-Moon L1", mu_em, collinear(mu_em)[0], (5.1476, 2.9321, 2.3344, 2.2688, 3.586, 1.48)),
                         ("Earth-Moon L2", mu_em, collinear(mu_em)[1], (3.1904, 2.1587, 1.8626, 1.7862, 2.913, 2.01)),
                         ("Sun-Earth L1", mu_se, collinear(mu_se)[0], (4.0611, 2.5327, 2.0865, 2.0152, 3.229, 22.95)),
                         ("Sun-Earth L2", mu_se, collinear(mu_se)[1], (3.9405, 2.4843, 2.0570, 1.9851, 3.187, 23.4))):
    c2, lam, nu, nz, k = collinear_linear(mu, xp)
    T_u = units("Earth-Moon" if mu is mu_em else "Sun-(Earth+Moon)")[2] / DAY
    check(f"{name}: c2", c2, st[0], 5e-5); check("  lambda", lam, st[1], 5e-5); check("  nu", nu, st[2], 5e-5)
    check("  nu_z = sqrt(c2)", nz, st[3], 5e-5); check("  axis ratio k", k, st[4], 5e-4)
    check("  e-folding time [d]", T_u / lam, st[5], 0.005)
    ev = linear_eigenvalues([xp, 0, 0], mu)
    check("  6x6 eigenvalues contain +-lambda", min(abs(abs(ev.real).max() - lam), 1), 0, 1e-8)
    truth("  and +-i nu, +-i nu_z", np.allclose(np.sort(np.abs(ev.imag))[::2][-2:], sorted([nu, nz]), atol=1e-8))
    H = hessian([xp, 0, 0], mu)
    check("  Omega_xx = 1 + 2 c2", H[0, 0], 1 + 2 * c2, 1e-9); check("  Omega_yy = 1 - c2", H[1, 1], 1 - c2, 1e-9)
    check("  Omega_zz = -c2", H[2, 2], -c2, 1e-9)
    check("  characteristic polynomial: lambda^4 + (2 - c2) lambda^2 + (1 + c2 - 2 c2^2)", lam ** 4 + (2 - c2) * lam ** 2 + (1 + c2 - 2 * c2 ** 2), 0, 1e-8)
c2s = c2_of(mu_se, collinear(mu_se)[0])
check("Sun-Earth L1: 1 + c2 - 2 c2^2", 1 + c2s - 2 * c2s ** 2, -27.92, 0.005)
check("  lambda^2 (positive root)", collinear_linear(mu_se, collinear(mu_se)[0])[1] ** 2, 6.4146, 5e-4)
check("  -nu^2", -collinear_linear(mu_se, collinear(mu_se)[0])[2] ** 2, -4.3535, 5e-4)
check("  linear in-plane period [d]", 2 * np.pi / collinear_linear(mu_se, collinear(mu_se)[0])[2] * T_se / DAY, 175.1, 0.05)
check("  linear out-of-plane period [d]", 2 * np.pi / collinear_linear(mu_se, collinear(mu_se)[0])[3] * T_se / DAY, 181.2, 0.05)
check("  in-plane period in time units", 2 * np.pi / collinear_linear(mu_se, collinear(mu_se)[0])[2], 3.0114, 5e-5)
check("  out-of-plane period in time units", 2 * np.pi / collinear_linear(mu_se, collinear(mu_se)[0])[3], 3.1179, 5e-5)
check("  e^{91/22.95}", np.exp(91 / 22.95), 53, 0.5)
check("ISRO target axis ratio 663200/209200", 663200 / 209200, 3.17, 0.005)
check("Earth-Moon L3 e-folding time [d]", T_em / DAY / abs(linear_eigenvalues([collinear(mu_em)[2], 0, 0], mu_em).real).max(), 24.4, 0.1)

print("Triangular points")
for mu in (0.1, mu_em, 0.3):
    L4 = triangular(mu)[0]; H = hessian(L4, mu)
    check(f"mu = {mu:.4f}: Omega_xx", H[0, 0], 0.75, 1e-12); check("  Omega_yy", H[1, 1], 2.25, 1e-12)
    check("  Omega_xy", H[0, 1], 3 * np.sqrt(3) / 4 * (1 - 2 * mu), 1e-12)
    check("  Omega_zz", H[2, 2], -1.0, 1e-12)
    ev = linear_eigenvalues(L4, mu)
    pol = lambda lam_: lam_ ** 4 + lam_ ** 2 + 27 / 4 * mu * (1 - mu)
    # the six eigenvalues are the four roots of pol and +-i (the vertical oscillation)
    rest = [l for l in ev if abs(abs(l.imag) - 1) > 1e-9 or abs(l.real) > 1e-9]
    check("  residual of the characteristic polynomial at the planar eigenvalues", max(abs(pol(l)) for l in rest), 0, 1e-8)
mu_R = (1 - np.sqrt(23 / 27)) / 2
check("Routh's mu", mu_R, 0.0385209, 5e-8)
check("  27 mu (1 - mu) = 1 there", 27 * mu_R * (1 - mu_R), 1.0, 1e-12)
truth("stable below Routh, unstable above (real parts of the six eigenvalues)", max(abs(linear_eigenvalues(triangular(0.035)[0], 0.035).real)) < 1e-6 and max(abs(linear_eigenvalues(triangular(0.045)[0], 0.045).real)) > 1e-3)
d = np.sqrt(1 - 27 * mu_em * (1 - mu_em))
w1, w2 = np.sqrt((1 + d) / 2), np.sqrt((1 - d) / 2)
check("Earth-Moon omega_1", w1, 0.9545, 5e-5); check("  omega_2", w2, 0.2982, 5e-5)
check("  periods [d]", 2 * np.pi / w1 * T_em / DAY, 28.6, 0.05); check("  and", 2 * np.pi / w2 * T_em / DAY, 91.5, 0.05)
truth("the 6x6 imaginary parts at L4 are omega_1, omega_2 and 1", np.allclose(sorted(set(np.round(np.abs(linear_eigenvalues(triangular(mu_em)[0], mu_em).imag), 8))), sorted([w2, w1, 1.0]), atol=1e-7))
mj = 1 / (1047.35 + 1)
dj = np.sqrt(1 - 27 * mj * (1 - mj)); wj1, wj2 = np.sqrt((1 + dj) / 2), np.sqrt((1 - dj) / 2)
check("Sun-Jupiter mu", mj, 9.54e-4, 5e-7)
check("  omega_1", wj1, 0.99676, 5e-6); check("  omega_2", wj2, 0.08046, 5e-6)
check("  periods in Jupiter years", 1 / wj1, 1.003, 5e-4); check("  and", 1 / wj2, 12.43, 0.005)
check("  slow period [years]", 11.862 / wj2, 147, 0.5)
check("27 * (1/2)(1/2) at mu = 1/2", 27 * 0.25, 6.75, 1e-12)
check("  |Im lambda^2| = sqrt(5.75)/2", np.sqrt(5.75) / 2, 1.1990, 5e-5)

print("Periodic orbits")
L1s = collinear(mu_se)[0]
x0, vy0, th = halo(mu_se, L1s - 0.0008, 120000.0 / L_se, 0.006)
check("halo period [d]", 2 * th * T_se / DAY, 177.87, 0.005)
ts, Yh = orbit_history(mu_se, [x0, 0, 120000.0 / L_se, 0, vy0, 0], 2 * th, 4000)
check("  A_x [km]", (Yh[0].max() - Yh[0].min()) / 2 * L_se, 208270, 5)
check("  A_y [km]", (Yh[1].max() - Yh[1].min()) / 2 * L_se, 667997, 5)
check("  closure: |state(T) - state(0)|", np.abs(Yh[:, -1] - Yh[:, 0]).max(), 0, 1e-8)
check("  Jacobi constant conserved (spread)", np.ptp([jacobi(Yh[:, k], mu_se) for k in range(0, 4000, 100)]), 0, 1e-11)
check("  period lies between the two linear ones [d]", float(175.1 < 2 * th * T_se / DAY < 181.2), 1.0, 0)
c2, lam, nu, nz, k = collinear_linear(mu_em, L1e)
per = []
for amp in (0.004, 0.044):
    pass
amps = np.arange(0.004, 0.0441, 0.004)
out = []
for amp in amps:
    if not out:
        guess = k * nu * amp
    elif len(out) == 1:
        guess = out[-1][1] * amp / out[-1][0]
    else:
        guess = out[-1][1] + (out[-1][1] - out[-2][1]) / (out[-1][0] - out[-2][0]) * (amp - out[-1][0])
    vy, th_l = lyapunov(mu_em, L1e - amp, guess)
    out.append((amp, vy, th_l))
check("Lyapunov, Earth-Moon L1, amplitude 0.004: period [d]", 2 * out[0][2] * T_em / DAY, 11.70, 0.005)
check("  amplitude 0.044: period [d]", 2 * out[-1][2] * T_em / DAY, 15.5, 0.05)
truth("  the period grows with the amplitude", all(b[2] > a[2] for a, b in zip(out[:-1], out[1:])))
s_end, t_end = _next_crossing([L1e - 0.044, 0, 0, 0, out[-1][1], 0], mu_em)
check("  perpendicular crossing: vx at the half period", s_end[3], 0, 1e-9)
check("Lyapunov at small amplitude against 2 pi / nu [d]", 2 * np.pi / nu * T_em / DAY, 11.70 if False else 2 * out[0][2] * T_em / DAY, 0.05)

print("Hill sphere and sphere of influence")
mE = 398600.4418 / 1.32712440018e11
check("Earth Hill radius [km]", 1.495978707e8 * (mE / 3) ** (1 / 3), 1.4966e6, 100)
check("Earth SOI [km]", 1.495978707e8 * mE ** 0.4, 9.25e5, 5e3)
check("  ratio", (mE / 3) ** (1 / 3) / mE ** 0.4, 1.6, 0.03)
check("Moon at the fraction of the Earth's Hill radius", 384400 / (1.495978707e8 * (mE / 3) ** (1 / 3)), 0.26, 0.005)
mM = 4902.800066 / 398600.4418
check("Moon: Hill radius [km]", 384400 * (mM / 3) ** (1 / 3), 61500, 100)
check("Moon: SOI [km]", 384400 * mM ** 0.4, 66200, 100)
check("crossover mass ratio where Hill = SOI", (3 ** (-1 / 3)) ** 15, 0.0041, 5e-5)

print("Exercises")
check("Ex 2: collinear L1 at mu = 1/2", collinear(0.5)[0], 0.0, 1e-12)
check("  triangular point y", triangular(0.5)[0][1], np.sqrt(3) / 2, 1e-12)
check("  triangular point x", triangular(0.5)[0][0], 0.0, 1e-12)
check("  L2 and L3 symmetric", collinear(0.5)[1] + collinear(0.5)[2], 0.0, 1e-9)
rl = 1737.4 + 100; MUm = 4902.800066
for th_d in (0, 90, 180, 270):
    th_ = np.radians(th_d)
    rb = np.array([1 - mu_em + rl / L_em * np.cos(th_), rl / L_em * np.sin(th_)])
    v_in = np.array([0, n_em * (1 - mu_em) * L_em]) + np.sqrt(MUm / rl) * np.array([-np.sin(th_), np.cos(th_)])
    v_rot = (v_in - n_em * L_em * np.array([-rb[1], rb[0]])) / V_em
    check(f"Ex 4: C of the 100 km lunar orbit at {th_d} deg", jacobi([rb[0], rb[1], 0, v_rot[0], v_rot[1], 0], mu_em), 5.5088, 1e-4)
check("  C - C_L1", 5.5088 - Cl1, 2.3, 0.05)

print(f"\n{sum(RESULTS)} PASS, {len(RESULTS) - sum(RESULTS)} FAIL")