"""Checks every derivation step and every number quoted in universal-variables.qmd.

Symbolic checks (sympy) for the derivations; numerical checks against direct integration of the
two-body equation and against the Lambert note for the worked examples.
Run from the sfm-slides directory:   python3 codes/universal-variables-verify.py
"""
import sys
import numpy as np
import sympy as sp
from scipy.integrate import solve_ivp
sys.path.insert(0, "codes")
from _uv import stumpff, kepler_uv, fg, propagate, lambert_uv, lambert_geometry, coe, MU_E

RESULTS = []
mu = MU_E


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


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


def integrate(r0, v0, dt):
    def rhs(t, y):
        r = y[:3]
        return np.r_[y[3:], -mu * r / np.linalg.norm(r) ** 3]
    s = solve_ivp(rhs, [0, dt], np.r_[r0, v0], rtol=1e-13, atol=1e-10, method="DOP853")
    return s.y[:3, -1], s.y[3:, -1]


# ------------------------------------------------------------------ 2. the radial equation becomes linear
r, rd, m, al, h = sp.symbols("r rdot mu alpha h", positive=True)
rd2 = 2 * m / r - m * al - h ** 2 / r ** 2                 # energy: rdot^2 = v^2 - h^2/r^2
rdd = -m / r ** 2 + h ** 2 / r ** 3                        # radial equation
truth("rdot^2 + r rddot = mu/r - mu alpha", sp.simplify(rd2 + r * rdd - (m / r - m * al)) == 0)
truth("r'' = (r/mu)(rdot^2 + r rddot) = 1 - alpha r", sp.simplify(r / m * (rd2 + r * rdd) - (1 - al * r)) == 0)

# solution of r'' + alpha r = 1 with r(0) = r0, r'(0) = sigma0, in trigonometric and Stumpff form
x, r0, s0, a = sp.symbols("chi r0 sigma0 alpha", positive=True)
w = sp.sqrt(a)
r_trig = 1 / a + (r0 - 1 / a) * sp.cos(w * x) + s0 / w * sp.sin(w * x)
truth("r(chi) satisfies r'' + alpha r = 1", sp.simplify(sp.diff(r_trig, x, 2) + a * r_trig - 1) == 0)
truth("r(0) = r0, r'(0) = sigma0", sp.simplify(r_trig.subs(x, 0) - r0) == 0 and sp.simplify(sp.diff(r_trig, x).subs(x, 0) - s0) == 0)
zz = a * x ** 2
Cs = (1 - sp.cos(sp.sqrt(zz))) / zz
Ss = (sp.sqrt(zz) - sp.sin(sp.sqrt(zz))) / sp.sqrt(zz) ** 3
r_st = x ** 2 * Cs + s0 * x * (1 - zz * Ss) + r0 * (1 - zz * Cs)
truth("Stumpff form of r(chi) equals the trigonometric form", sp.simplify(r_st - r_trig) == 0)
truth("cos(sqrt(alpha) chi) = 1 - zC", sp.simplify(sp.cos(w * x) - (1 - zz * Cs)) == 0)
truth("sin(sqrt(alpha) chi)/sqrt(alpha) = chi (1 - zS)", sp.simplify(sp.sin(w * x) / w - x * (1 - zz * Ss)) == 0)
truth("(1 - cos(sqrt(alpha) chi))/alpha = chi^2 C", sp.simplify((1 - sp.cos(w * x)) / a - x ** 2 * Cs) == 0)

# integral of r gives the universal Kepler equation
I = sp.integrate(r_trig, (x, 0, x))
kep = r0 * x + (1 - a * r0) * x ** 3 * Ss + s0 * x ** 2 * Cs
truth("integral of r dchi = r0 chi + (1 - alpha r0) chi^3 S + sigma0 chi^2 C", sp.simplify(I - kep) == 0)
truth("d(universal Kepler)/dchi = r", sp.simplify(sp.diff(kep, x) - r_trig) == 0)

# derivative identities of the three basis functions
truth("d/dchi [chi(1 - zS)] = 1 - zC", sp.simplify(sp.diff(x * (1 - zz * Ss), x) - (1 - zz * Cs)) == 0)
truth("d/dchi [chi^2 C] = chi(1 - zS)", sp.simplify(sp.diff(x ** 2 * Cs, x) - x * (1 - zz * Ss)) == 0)
truth("d/dchi [chi^3 S] = chi^2 C", sp.simplify(sp.diff(x ** 3 * Ss, x) - x ** 2 * Cs) == 0)

# series
zs = sp.symbols("z")
cser = sp.series((1 - sp.cos(sp.sqrt(zs))) / zs, zs, 0, 3).removeO()
sser = sp.series((sp.sqrt(zs) - sp.sin(sp.sqrt(zs))) / zs ** sp.Rational(3, 2), zs, 0, 3).removeO()
truth("C(z) = 1/2 - z/24 + z^2/720 - ...", sp.simplify(cser - (sp.Rational(1, 2) - zs / 24 + zs ** 2 / 720)) == 0)
truth("S(z) = 1/6 - z/120 + z^2/5040 - ...", sp.simplify(sser - (sp.Rational(1, 6) - zs / 120 + zs ** 2 / 5040)) == 0)
for zv in (-25.0, -0.5, -1e-3, 0.0, 1e-3, 0.5, 25.0):
    C, S = stumpff(zv)
    Cr = float(sp.N(cser.subs(zs, 0) if zv == 0 else ((1 - sp.cos(sp.sqrt(sp.Float(zv, 40)))) / sp.Float(zv, 40)).evalf(30)))
    Sr = float(sp.N(sp.Rational(1, 6) if zv == 0 else ((sp.sqrt(sp.Float(zv, 40)) - sp.sin(sp.sqrt(sp.Float(zv, 40)))) / sp.sqrt(sp.Float(zv, 40)) ** 3).evalf(30)))
    check(f"Stumpff C at z={zv} (30-digit reference)", C, Cr, 1e-15)
    check(f"Stumpff S at z={zv} (30-digit reference)", S, Sr, 1e-15)
C0, S0 = stumpff(0.0)
check("C(0)", C0, 0.5, 0)
check("S(0)", S0, 1 / 6, 1e-17)
check("C(4 pi^2) = 0 (a full turn)", stumpff(4 * np.pi ** 2)[0], 0.0, 1e-16)

# the third-order vector equation, and F, G
F_ex = 1 - x ** 2 * Cs / r0
truth("F(chi) = 1 - chi^2 C/r0 satisfies F''' + alpha F' = 0",
      sp.simplify(sp.diff(F_ex, x, 3) + a * sp.diff(F_ex, x)) == 0)
basis = [sp.Integer(1), x * (1 - zz * Ss), x ** 2 * Cs]
truth("1, chi(1-zS), chi^2 C all satisfy u''' + alpha u' = 0",
      all(sp.simplify(sp.diff(b, x, 3) + a * sp.diff(b, x)) == 0 for b in basis))
vals = [[sp.limit(sp.diff(b, x, k), x, 0) for b in basis] for k in range(3)]
truth("basis values at chi = 0: identity matrix", vals == [[1, 0, 0], [0, 1, 0], [0, 0, 1]])
# third-order equation from r r'' - r' r' + r = 0 (vector, component-free check with r'' + alpha r = 1)
rv = sp.Function("R")(x)
rs = sp.Function("rho")(x)
eq2 = rs * sp.diff(rv, x, 2) - sp.diff(rs, x) * sp.diff(rv, x) + rv
d_eq2 = sp.expand(sp.diff(eq2, x))
target = rs * sp.diff(rv, x, 3) + (1 - sp.diff(rs, x, 2)) * sp.diff(rv, x)
truth("d/dchi [r R'' - r' R' + R] = r R''' + (1 - r'') R'", sp.simplify(d_eq2 - target) == 0)

# reduction to the periapsis form of the two-body note: r0 = r_p, sigma0 = 0
rp_ = sp.symbols("r_p", positive=True)
kep_p = kep.subs({r0: rp_, s0: 0})
truth("periapsis start gives r_p chi + (1 - r_p/a) chi^3 S (two-body note)",
      sp.simplify(kep_p - (rp_ * x + (1 - rp_ * a) * x ** 3 * Ss)) == 0)

# ------------------------------------------------------------------ meaning of chi on each conic
r0v = np.array([-2757.4, -6465.5, -1726.8])
v0v = np.array([7.4488, -4.3468, -4.4076])
dt1 = 3 * 3600.0
nr0 = np.linalg.norm(r0v)
vr0 = r0v @ v0v / nr0
alpha = 2 / nr0 - v0v @ v0v / mu
aa, ee, ii, Om, ww, f0 = coe(r0v, v0v)
check("E1 r0 [km]", nr0, 7237.94, 0.005)
check("E1 v_r0 [km/s]", vr0, 2.0967, 5e-5)
check("E1 v0 [km/s]", np.linalg.norm(v0v), 9.6854, 5e-5)
check("E1 alpha [1/km]", alpha * 1e5, 4.0983, 5e-5)
check("E1 a [km]", aa, 24400.5, 0.05)
check("E1 e", ee, 0.72001, 5e-6)
check("E1 f0 [deg]", f0, 30.00, 0.005)
chi1, hist1 = kepler_uv(nr0, vr0, alpha, dt1, history=True)
check("E1 chi_0 = sqrt(mu)|alpha|dt [km^1/2]", hist1[0], 279.443, 5e-4)
for k, stated in enumerate([279.443, 339.352, 334.779, 334.757, 334.757]):
    check(f"E1 Newton iterate {k}", hist1[k], stated, 5e-4)
check("E1 Newton steps to 1e-12", len(hist1) - 1, 5, 0)
F1, G1, Fd1, Gd1, chi_, _, rr1 = fg(r0v, v0v, dt1)
z1 = alpha * chi1 ** 2
check("E1 z", z1, 4.5926, 5e-5)
check("E1 C(z)", stumpff(z1)[0], 0.33565, 5e-6)
check("E1 S(z)", stumpff(z1)[1], 0.13232, 5e-6)
check("E1 F", F1, -4.19676, 5e-6)
check("E1 G [s]", G1, 2937.580, 5e-4)
check("E1 Fdot [1/s]", Fd1 * 1e4, -3.10843, 5e-6)
check("E1 Gdot", Gd1, -0.020700, 5e-7)
check("E1 r [km]", rr1, 36850.99, 0.005)
check("E1 F Gdot - Fdot G", F1 * Gd1 - Fd1 * G1, 1.0, 1e-12)
r1v, v1v = propagate(r0v, v0v, dt1)
ri, vi = integrate(r0v, v0v, dt1)
for k, s in enumerate([33453.581, 14365.062, -5700.716]):
    check(f"E1 r[{k}] [km]", r1v[k], s, 5e-4)
for k, s in enumerate([0.70292, 2.09973, 0.62800]):
    check(f"E1 v[{k}] [km/s]", v1v[k], s, 5e-6)
check("E1 |r - numerical integration| [km]", np.linalg.norm(r1v - ri), 0, 1e-6)
check("E1 |v - numerical integration| [km/s]", np.linalg.norm(v1v - vi), 0, 1e-9)
check("E1 f after 3 h [deg]", coe(r1v, v1v)[5], 161.08, 0.005)
# chi = sqrt(a) dE
f_a, f_b = np.radians(f0), np.radians(coe(r1v, v1v)[5])
Ea = 2 * np.arctan(np.sqrt((1 - ee) / (1 + ee)) * np.tan(f_a / 2))
Eb = 2 * np.arctan(np.sqrt((1 - ee) / (1 + ee)) * np.tan(f_b / 2))
check("E1 chi / sqrt(a) = Delta E [rad]", chi1 / np.sqrt(aa), Eb - Ea, 1e-9)
check("E1 Delta E [deg]", np.degrees(Eb - Ea), 122.79, 0.005)

# hyperbola: Example 2, the departure hyperbola of the patched-conics note
rp2, vinf = 6678.137, 2.9446935
vp2 = np.sqrt(vinf ** 2 + 2 * mu / rp2)
r02, v02 = np.array([rp2, 0, 0.0]), np.array([0, vp2, 0.0])
dt2 = 86400.0
alpha2 = 2 / rp2 - vp2 ** 2 / mu
check("E2 v_p [km/s]", vp2, 11.316, 5e-4)
check("E2 alpha [1/km]", alpha2 * 1e5, -2.1754, 5e-5)
check("E2 a [km]", 1 / alpha2, -45968, 1)
chi2, hist2 = kepler_uv(rp2, 0.0, alpha2, dt2, history=True)
check("E2 chi [km^1/2]", chi2, 571.73, 0.005)
check("E2 z", alpha2 * chi2 ** 2, -7.1110, 5e-4)
check("E2 Newton steps", len(hist2) - 1, 6, 0)
r2v, v2v = propagate(r02, v02, dt2)
ri2, vi2 = integrate(r02, v02, dt2)
check("E2 r after 1 day [km]", np.linalg.norm(r2v), 334692, 1)
check("E2 v after 1 day [km/s]", np.linalg.norm(v2v), 3.3246, 5e-4)
check("E2 |r - integration| [km]", np.linalg.norm(r2v - ri2), 0, 1e-5)
ecc2 = 1 + rp2 * vinf ** 2 / mu
Hb = np.arccosh((1 + np.linalg.norm(r2v) / 45960.0) / ecc2) if False else None
a2 = 1 / alpha2
fb2 = np.radians(coe(r2v, v2v)[5])
H2 = 2 * np.arctanh(np.sqrt((ecc2 - 1) / (ecc2 + 1)) * np.tan(fb2 / 2))
check("E2 chi / sqrt(-a) = Delta H", chi2 / np.sqrt(-a2), H2, 1e-8)
check("E2 Delta H", H2, 2.6666, 5e-4)
check("E2 true anomaly after 1 day [deg]", np.degrees(fb2), 146.70, 0.005)
check("E2 f_inf [deg] (approached, not reached)", np.degrees(np.arccos(-1 / ecc2)), 150.83, 0.005)

# parabola: chi = sqrt(p) (tan f/2 - tan f0/2)
rpp = 7000.0
vpp = np.sqrt(2 * mu / rpp)
rP, vP = propagate(np.array([rpp, 0, 0.0]), np.array([0, vpp, 0.0]), 3000.0)
nrp = np.linalg.norm(rP)
alphaP = 2 / rpp - vpp ** 2 / mu
check("parabola alpha", alphaP, 0.0, 1e-18)
chiP = kepler_uv(rpp, 0.0, 0.0, 3000.0)
fP = np.arctan2(rP[1], rP[0])
check("parabola chi = sqrt(p) tan(f/2)", chiP, np.sqrt(2 * rpp) * np.tan(fP / 2), 1e-9)
truth("parabola: sqrt(mu) t = r_p chi + chi^3/6 (Barker)", abs(np.sqrt(mu) * 3000 - (rpp * chiP + chiP ** 3 / 6)) < 1e-6)

# ------------------------------------------------------------------ precision near e = 1 (figure script reference)
import importlib.util
spec = importlib.util.spec_from_file_location("figuv", "codes/fig-uv.py")
figuv = importlib.util.module_from_spec(spec)
import matplotlib
matplotlib.use("Agg")
spec.loader.exec_module(figuv)
eu, ec = figuv.precision_case(1 - 1e-8)
check("precision e=1-1e-8: universal error [m]", eu * 1e3, 0.0, 1e-7)
check("precision e=1-1e-8: classical error [m] (order)", np.log10(ec * 1e3), np.log10(0.45), 0.6)
eu, ec = figuv.precision_case(1 - 1e-2)
check("precision e=0.99: classical error [m] (order)", np.log10(ec * 1e3), np.log10(6e-8), 0.6)

# ------------------------------------------------------------------ Lambert in universal variables
th1, th2, rr1_, rr2_ = sp.symbols("Delta_theta theta2 r_1 r_2", positive=True)
A_def = sp.sin(th1) * sp.sqrt(rr1_ * rr2_ / (1 - sp.cos(th1)))
A_alt = sp.sqrt(2 * rr1_ * rr2_) * sp.cos(th1 / 2)
truth("A = sqrt(2 r1 r2) cos(dth/2)  (checked at 9 angles)",
      all(abs(float((A_def - A_alt).subs({th1: t, rr1_: 6678.0, rr2_: 42164.0}))) < 1e-8
          for t in np.linspace(0.2, 6.1, 9)))
truth("sqrt(r1 r2 (1-cos))/tan(dth/2) = A  (checked at 9 angles)",
      all(abs(float((sp.sqrt(rr1_ * rr2_ * (1 - sp.cos(th1))) / sp.tan(th1 / 2) - A_def).subs({th1: t, rr1_: 6678.0, rr2_: 42164.0}))) < 1e-7
          for t in np.linspace(0.2, 6.1, 9)))
r1L = np.array([6678.137, 0, 0.0])
thL = np.radians(175)
r2L = 42164 * np.array([np.cos(thL), np.sin(thL), 0.0])
dth, A, n1, n2 = lambert_geometry(r1L, r2L)
check("E3 A [km]", A, 1035.126, 5e-4)
outL, histL = lambert_uv(r1L, r2L, 14400.0, history=True)
check("E3 z", outL["z"], 6.85244, 5e-6)
check("E3 Newton steps", len(histL) - 1, 7, 0)
check("E3 y [km]", outL["y"], 48463.06, 0.005)
check("E3 chi = sqrt(y/C) [km^1/2]", np.sqrt(outL["y"] / stumpff(outL["z"])[0]), 421.88, 0.005)
check("E3 F", outL["F"], -6.25697, 5e-5)
check("E3 G [s]", outL["G"], 360.936, 5e-3)
check("E3 Gdot", outL["Gd"], -0.14939, 5e-5)
check("E3 p = r1 r2 (1-cos)/y [km]", n1 * n2 * (1 - np.cos(dth)) / outL["y"], 11598.163, 5e-3)
for k, s in enumerate([-0.60573, 10.18142]):
    check(f"E3 v1[{k}] (Lambert note Example 1)", outL["v1"][k], s, 5e-5)
for k, s in enumerate([-1.11668, -1.52104]):
    check(f"E3 v2[{k}] (Lambert note Example 1)", outL["v2"][k], s, 5e-5)
rL, vL = propagate(r1L, outL["v1"], 14400.0)
check("E3 propagating v1 for 4 h lands on r2 [km]", np.linalg.norm(rL - r2L), 0, 1e-6)
out2h = lambert_uv(r1L, r2L, 7200.0)
check("E3' 2 h: z < 0 (hyperbola)", out2h["z"], -1.4807, 5e-4)
check("E3' 2 h: e from v1", coe(r1L, out2h["v1"])[1], 1.0, 0.5)
truth("E3' 2 h: hyperbola, e > 1", coe(r1L, out2h["v1"])[1] > 1)
check("E3' 2 h: v1 y-component [km/s]", out2h["v1"][1], 10.32782, 5e-5)
# monotonic time of flight in z
from _uv import lambert_y, lambert_tof
zs_ = np.linspace(-60, 4 * np.pi ** 2 - 0.01, 4000)
tv = np.array([lambert_tof(v, n1, n2, A) for v in zs_])
ok = ~np.isnan(tv)
truth("E3 time of flight increases with z wherever y > 0", np.all(np.diff(tv[ok]) > 0))
zcut = zs_[~ok].max() if (~ok).any() else None
check("E3 y = 0 at z [approx]", max(z for z in np.linspace(-80, 0, 80001) if lambert_y(z, n1, n2, A) < 0), -70.6, 0.05)

# ------------------------------------------------------------------ numbers quoted in the text and boxes
check("E1 sqrt(z) = Delta E [rad]", np.sqrt(z1), 2.143, 5e-4)
check("E3 y(0) [km]", lambert_y(0.0, n1, n2, A), 47378.25, 0.005)
check("E3 parabolic time at z = 0 [h] (Lambert note: 2.238 h)", lambert_tof(0.0, n1, n2, A) / 3600, 2.238, 5e-4)
for k, (zs_k, ys_k, ts_k) in enumerate([(0, 47378.25, 2.238033), (10.18475, 48878.56, 5.553954), (7.49666, 48548.73, 4.250411),
                                         (6.87788, 48466.49, 4.009509), (6.85248, 48463.06, 4.000015), (6.85244, 48463.06, 4.000000)]):
    check(f"E3 table step {k} z", histL[k], zs_k, 5e-5)
    check(f"E3 table step {k} y", lambert_y(histL[k], n1, n2, A), ys_k, 0.005)
    check(f"E3 table step {k} t", lambert_tof(histL[k], n1, n2, A) / 3600, ts_k, 5e-7)
check("E3' 2 h: e", coe(r1L, out2h["v1"])[1], 1.1009, 5e-5)
eS = 1 + rp2 * vinf ** 2 / mu
coshH = (1 - 925000.0 * alpha2) / eS
HS = np.arccosh(coshH)
chiS = np.sqrt(-a2) * HS
zS = alpha2 * chiS ** 2
CSs = stumpff(zS)
tS = ((1 - alpha2 * rp2) * chiS ** 3 * CSs[1] + rp2 * chiS) / np.sqrt(mu)
check("task: e of the departure hyperbola", eS, 1.1453, 5e-5)
check("task: cosh H at the SOI", coshH, 18.44, 0.005)
check("task: H at the SOI", HS, 3.607, 5e-4)
check("task: chi at the SOI", chiS, 773.4, 0.05)
check("task: z at the SOI", zS, -13.01, 0.005)
check("task: time to the SOI [s]", tS / 1e5, 2.73, 0.005)
check("task: time to the SOI [d]", tS / 86400, 3.16, 0.005)
check("E2 F", fg(r02, v02, dt2)[0], -41.887, 5e-4)
check("E2 G [s]", fg(r02, v02, dt2)[1], 16240.18, 0.005)
check("E2 speed from energy at r [km/s]", np.sqrt(vinf ** 2 + 2 * mu / np.linalg.norm(r2v)), 3.3246, 5e-4)
check("E2 chi_0 (logarithmic start)", hist2[0], 486.38, 0.005)
# elliptic reduction of F and G (Example 1 numbers)
dE = np.sqrt(z1)
check("E1 F = 1 - a/r0 (1 - cos dE)", 1 - aa / nr0 * (1 - np.cos(dE)), F1, 1e-9)
check("E1 G = dt - sqrt(a^3/mu)(dE - sin dE)", dt1 - np.sqrt(aa ** 3 / mu) * (dE - np.sin(dE)), G1, 1e-6)

npass = sum(ok for _, ok in RESULTS)
print(f"\n{npass} PASS, {len(RESULTS) - npass} FAIL")