"""Checks every number and identity quoted in impulsive-model.qmd. Prints FAIL on a mismatch.
Run from anywhere:   python3 codes/impulsive-model-verify.py"""
import numpy as np
import sympy as sp

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, g0 = 398600.4418, 6378.137, 9.80665e-3
vis = lambda r, a: np.sqrt(mu*(2/r - 1/a))

# Mangalyaan: orbit from codes/_imcommon.py (MOM_EARTH); 1329 s and 648 m/s from indian-missions-problems.qmd; 616 m/s from escape.qmd
rp, ra = Re + 264.1, Re + 192874
a = (rp + ra)/2; T = 2*np.pi*np.sqrt(a**3/mu)
ok("period h", T/3600, 91, 0.5); ok("t_burn/T percent", 100*1329/T, 0.4, 0.05)
vp = vis(rp, a); ok("speed at perigee", vp, 10.8, 0.05); ok("distance in the burn, km", vp*1329, 14000, 500)
ok("impulsive m/s", 1e3*(np.sqrt(3.129**2 + 2*mu/rp) - vp), 616, 0.5); ok("about 5 percent less", 100*(648 - 616)/648, 5, 0.1)

# the rocket equation, by integrating the momentum balance
m, v, dm, dv, ve_ = sp.symbols("m v dm dv v_e")
ok("expansion", sp.expand((m + dm)*(v + dv) + (-dm)*(v - ve_) - m*v - (m*dv + ve_*dm + dm*dv)), 0, 0)
ve = 315*g0
ok("v_e", ve, 3.0891, 5e-5); ok("v_e, 3 places", ve, 3.089, 5e-4)
for d, L in ((1, 72.3), (2, 52.3), (4, 27.4), (8, 7.5)):
    ok(f"mass left at {d} km/s", 100*np.exp(-d/ve), L, 0.05)
ok("doubling squares", np.exp(-4/ve), np.exp(-2/ve)**2, 1e-12)
# a numerical burn: small steps agree with the formula
M, V = 1.0, 0.0
for _ in range(20000):
    dM = -M*1e-4; V += -ve*dM/M; M += dM
ok("stepwise mass balance", V, ve*np.log(1/M), 1e-3)

# the example
D = 3.8926
ok("exponent", D/ve, 1.2601, 5e-5); ok("propellant", 1 - np.exp(-D/ve), 0.716, 5e-4); ok("left percent", 100*np.exp(-D/ve), 28.36, 5e-3)
ok("after saving", 100*np.exp(-(D - 0.05)/ve), 28.83, 5e-3); ok("kg", 3000*(np.exp(-(D - 0.05)/ve) - np.exp(-D/ve)), 14, 0.5)
ok("gain percent", 100*(np.exp(0.05/ve) - 1), 1.6, 0.05); ok("saving in dv percent", 100*0.05/D, 1.3, 0.05)

# what an impulse does
rng = np.random.default_rng(4)
for _ in range(20):
    r = rng.normal(size=3)*7000; v_ = rng.normal(size=3)*5; d = rng.normal(size=3)
    e0 = v_ @ v_/2 - mu/np.linalg.norm(r); e1 = (v_ + d) @ (v_ + d)/2 - mu/np.linalg.norm(r)
    ok("energy change", e1 - e0, v_ @ d + d @ d/2, 1e-9)
    ok("h change", np.linalg.norm(np.cross(r, v_ + d) - np.cross(r, v_) - np.cross(r, d)), 0, 1e-8)
    b = np.arccos(v_ @ d/np.linalg.norm(v_)/np.linalg.norm(d))
    ok("the cos(beta) form", np.linalg.norm(v_)*np.linalg.norm(d)*np.cos(b) + d @ d/2, e1 - e0, 1e-9)
# a radial burn keeps h and p; off the apsides it changes the energy at first order by v_r dv
r = np.array((8000.0, 0.0, 0.0)); v_ = np.array((1.5, 7.0, 0.0)); d = 0.01*r/np.linalg.norm(r)
ok("radial: h unchanged", np.linalg.norm(np.cross(r, v_ + d) - np.cross(r, v_)), 0, 1e-9)
ok("radial: first order term", v_ @ d, 1.5*0.01, 1e-12)
# a pure turn: the isosceles triangle, no energy change
for th in np.radians((5.0, 30.0, 90.0)):
    a1, a2 = np.array((7.0, 0.0, 0.0)), 7*np.array((np.cos(th), 0.0, np.sin(th)))
    ok("turn cost", np.linalg.norm(a2 - a1), 2*7*np.sin(th/2), 1e-12); ok("no energy", a2 @ a2 - a1 @ a1, 0, 1e-12)

# the example on an ellipse
a, e, D = 12000.0, 0.5, 0.1
vP, vA = np.sqrt(mu/a*(1 + e)/(1 - e)), np.sqrt(mu/a*(1 - e)/(1 + e))
ok("v_p", vP, 9.982, 5e-4); ok("v_a", vA, 3.327, 5e-4); ok("vis-viva agrees", vP, vis(a*(1 - e), a), 1e-12); ok("vis-viva agrees", vA, vis(a*(1 + e), a), 1e-12)
ok("energy at periapsis", vP*D + D*D/2, 1.003, 5e-4); ok("energy at apoapsis", vA*D + D*D/2, 0.338, 5e-4)
ok("nearly three times", (vP*D + D*D/2)/(vA*D + D*D/2), 3, 0.05); ok("speed ratio", vP/vA, 3, 1e-12)

# the exercise
ve = 320*g0
ok("left", 100*np.exp(-5/ve), 20.3, 0.05); ok("payload", 100*(np.exp(-5/ve) - 0.1), 10.3, 0.05)
s = (np.exp(-2.5/ve) - 0.1)**2; ok("two stages", 100*s, 12.3, 0.05); ok("more", 100*(s/(np.exp(-5/ve) - 0.1) - 1), 19, 0.5)

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