"""Checks every number and identity quoted in escape.qmd. Prints FAIL on a mismatch.
Run from anywhere:   python3 codes/escape-verify.py"""
import numpy as np
import sympy as sp
from scipy.integrate import quad
from scipy.optimize import brentq

fails = 0


def ok(name, x, y, tol):
    global fails
    good = abs(float(x) - float(y)) <= tol
    fails += not good
    print("ok  " if good else "FAIL", name, round(float(x), 6), y)


mu, Re = 398600.4418, 6378.137
rp = 6678.137
ve = 9.80665e-3*320
left = lambda dv: np.exp(-dv/ve)
vis = lambda r, a: np.sqrt(mu*(2/r - 1/a))
esc = lambda r: np.sqrt(2*mu/r)
dve = lambda r, a: esc(r) - vis(r, a)
half = lambda a: np.pi*np.sqrt(a**3/mu)
tpar = lambda q, R: (2/3)*np.sqrt((R - q)/(2*mu))*(R + 2*q)

# the form in x, and that it falls with x
a = 12000.0
for r in (8000.0, 15000.0):
    x = a/r
    ok("the form in x", dve(r, a), np.sqrt(mu/a)*(np.sqrt(2*x) - np.sqrt(2*x - 1)), 1e-12)
    ok("conjugate form", np.sqrt(2*x) - np.sqrt(2*x - 1), 1/(np.sqrt(2*x) + np.sqrt(2*x - 1)), 1e-12)
fx = [dve(a*(1 - .5**2)/(1 + .5*np.cos(f)), a) for f in np.linspace(0, np.pi, 200)]
ok("cost rises from periapsis to apoapsis", all(np.diff(fx) > 0), 1, 0)
# a burn not along the velocity is never cheaper
rng = np.random.default_rng(2); worst = 1e9
for _ in range(2000):
    v0 = rng.uniform(1, 10); ve_ = v0 + rng.uniform(0.1, 3); t = rng.uniform(0, np.pi)
    worst = min(worst, np.linalg.norm(ve_*np.array((np.cos(t), np.sin(t))) - (v0, 0)) - (ve_ - v0))
ok("the prograde burn is the smallest", worst >= -1e-12, 1, 0)

# v_p / v_esc
for e in (0.3, 0.9):
    ok("v_p/v_esc", vis(rp, rp/(1 - e))/esc(rp), np.sqrt((1 + e)/2), 1e-12)
ok("e = 0.9, percent", 100*np.sqrt(1.9/2), 97, 0.5)

# the table
for e, P, A, pct in ((0, 3.2001, 3.2001, 0), (0.3, 2.1172, 3.2742, 55), (0.5, 1.4638, 3.1540, 115), (0.7, 0.8527, 2.8122, 230), (0.9, 0.2766, 1.9461, 603)):
    a = rp/(1 - e); ra = a*(1 + e)
    ok(f"e={e} peri", dve(rp, a), P, 6e-5); ok(f"e={e} apo", dve(ra, a), A, 6e-5); ok(f"e={e} extra", 100*(dve(ra, a)/dve(rp, a) - 1), pct, 0.5)
a = rp/0.1; ok("seven times", dve(a*1.9, a)/dve(rp, a), 7, 0.1)
ok("mass left peri", 100*left(dve(rp, a)), 91.6, 0.05); ok("mass left apo", 100*left(dve(a*1.9, a)), 53.8, 0.05)

# the parabola
q, R_, r_ = sp.symbols("q R r", positive=True)
ok("v_r^2", sp.simplify(2/r_ - 2*q/r_**2 - (2/r_)*(1 - q/r_)), 0, 0)
for Q, RR in ((6678.0, 9.25e5), (1e5, 9.25e5)):
    I, _ = quad(lambda s: 1/(np.sqrt(2*mu/s)*np.sqrt(1 - Q/s)), Q, RR, limit=200)
    ok("the integral", I, tpar(Q, RR), 1e-2)
t_s = sp.Rational(2, 3)*sp.sqrt(R_ - q)*(R_ + 2*q)
ok("dt/dq", sp.simplify(sp.diff(t_s, q) - (R_ - 2*q)/sp.sqrt(R_ - q)), 0, 0)
R = 9.25e5
ok("R/2", R/2/1e5, 4.62, 0.005)

# the thresholds
ok("the printed R is the rounded sphere of influence, 1e5 km", 1.495978707e8*3.0034896e-6**0.4/1e5, 9.25, 0.01)
pen = lambda e: tpar(rp, R) + half(rp/(1 - e)) - tpar(rp*(1 + e)/(1 - e), R)
ok("first threshold", brentq(pen, 0.1, 0.3), 0.216, 5e-4); ok("second threshold", brentq(pen, 0.9, 0.949), 0.939, 6e-4)
ok("slower below", pen(0.1) > 0, 1, 0); ok("faster between", pen(0.5) < 0, 1, 0); ok("slower above", pen(0.95) > 0, 1, 0)

# the example
e = 0.9; a = rp/(1 - e); ra = a*(1 + e)
ok("a", a, 66781.37, 5e-3); ok("r_a", ra, 126884.6, 0.05)
ok("half period h", half(a)/3600, 23.9, 0.05); ok("climb h", tpar(rp, R)/3600, 186.5, 0.05); ok("sum h", (half(a) + tpar(rp, R))/3600, 210.4, 0.05)
ok("at once h", tpar(ra, R)/3600, 218.4, 0.05); ok("saves km/s", dve(ra, a) - dve(rp, a), 1.67, 5e-3); ok("sooner h", (tpar(ra, R) - half(a) - tpar(rp, R))/3600, 8, 0.1)
e = 0.95; a = rp/(1 - e); ra = a*(1 + e)
ok("e = 0.95 saves", dve(ra, a) - dve(rp, a), 1.34, 5e-3); ok("half period days", half(a)/86400, 2.8, 0.05); ok("later h", pen(0.95)/3600, 9.5, 0.05)

# Mangalyaan: the orbit is from codes/_imcommon.py (MOM_EARTH), the v_inf and the flown burn from indian-missions-problems.qmd
rP, rA = Re + 264.1, Re + 192874
a = (rP + rA)/2; e = (rA - rP)/(rA + rP); vp, va = vis(rP, a), vis(rA, a)
ok("r_p", rP, 6642.2, 0.05); ok("r_a", rA, 199252.1, 0.05); ok("a", a, 102947.2, 0.05); ok("e", e, 0.9355, 5e-5); ok("v_p", vp, 10.777, 5e-4); ok("v_a", va, 0.359, 5e-4)
vinf = 3.129
ok("esc at perigee", esc(rP), 10.955, 5e-4); ok("dv", esc(rP) - vp, 0.178, 5e-4); ok("esc at apogee", esc(rA), 2.000, 5e-4); ok("dv", esc(rA) - va, 1.641, 5e-4)
hp, ha = np.sqrt(vinf**2 + esc(rP)**2), np.sqrt(vinf**2 + esc(rA)**2)
ok("needed at perigee", hp, 11.393, 5e-4); ok("dv", hp - vp, 0.616, 5e-4); ok("needed at apogee", ha, 3.714, 5e-4); ok("dv", ha - va, 3.354, 5e-4)
ok("two thirds used, I_sp = 315 s", 1 - np.exp(-(ha - va)/(315*9.80665e-3)), 2/3, 0.02); ok("half period h", half(a)/3600, 46, 0.5)

print("\nfailures:", fails)
raise SystemExit(1 if fails else 0)