"""Checks every number and identity quoted in bielliptic-transfer.qmd. Prints FAIL on a mismatch.
Run from anywhere:   python3 codes/bielliptic-transfer-verify.py"""
import numpy as np
import sympy as sp
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, r1 = 398600.4418, 6678.137
vc1 = np.sqrt(mu/r1)
ve = 9.80665e-3*320
left = lambda dv: np.exp(-dv/ve)
vis = lambda r, a: np.sqrt(mu*(2/r - 1/a))
half = lambda a: np.pi*np.sqrt(a**3/mu)
g = lambda R: np.sqrt(2*R/(1 + R)) - 1 + (1 - np.sqrt(2/(1 + R)))/np.sqrt(R)


def bi(R, S):
    """The three burns in units of v_c1, positive sizes, from the formulas of the note; S = r_b/r_1 > R."""
    return (np.sqrt(2*S/(1 + S)) - 1, np.sqrt(2/S)*(np.sqrt(R/(R + S)) - np.sqrt(1/(1 + S))), np.sqrt(2*S/(R*(R + S))) - 1/np.sqrt(R))


# the speed formula on an ellipse with apsides r and r'
for r, rr in ((7000., 50000.), (42000., 9000.)):
    ok("v from apsides", np.sqrt(2*mu*rr/(r*(r + rr))), vis(r, (r + rr)/2), 1e-10)
ok("R* = R is Hohmann", sum(bi(15, 15)), g(15), 1e-12); ok("no third burn", bi(15, 15)[2], 0, 1e-12)
ok("second burn vanishes as r_b grows", bi(15, 1e12)[1], 0, 1e-4)

# the example
R, rb = 15.0, 400000.0
r2 = R*r1
ok("r2", r2, 100172.06, 6e-3)
a1, a2 = (r1 + rb)/2, (r2 + rb)/2
ok("a_t1", a1, 203339, 0.5); ok("a_t2", a2, 250086, 0.5)
rows = ((r1, vc1, vis(r1, a1), 7.72576, 10.83580, 3.11004), (rb, vis(rb, a1), vis(rb, a2), 0.18091, 0.63178, 0.45088), (r2, vis(r2, a2), np.sqrt(mu/r2), 2.52279, 1.99478, 0.52801))
tot = 0
for r, b, c, B, C, D in rows:
    ok("before", b, B, 6e-6); ok("after", c, C, 6e-6); ok("dv", abs(c - b), D, 6e-6); tot += abs(c - b)
ok("total", tot, 4.08892, 6e-6); ok("formula agrees", sum(bi(R, rb/r1))*vc1, tot, 1e-9)
H = g(R)*vc1; ok("Hohmann", H, 4.14269, 6e-6); ok("saving m/s", 1e3*(H - tot), 53.8, 0.05); ok("saving %", 100*(H - tot)/H, 1.3, 0.05)
tb = half(a1) + half(a2); ok("time h", tb/3600, 299.6, 0.05); ok("time days", tb/86400, 12.5, 0.05); ok("Hohmann h", half((r1 + r2)/2)/3600, 17.1, 0.05)
ok("periapsis raised, km", r2 - r1, 93000, 1000); ok("speed at r_b, m/s", vis(rb, a1)*1e3, 181, 0.5); ok("burn 2", abs(rows[1][2] - rows[1][1]), 0.45, 5e-3)
ok("mass bi %", 100*left(tot), 27.2, 0.05); ok("mass H %", 100*left(H), 26.7, 0.05); ok("more mass %", 100*(left(tot)/left(H) - 1), 1.7, 0.05)
ok("extra days", (tb - half((r1 + r2)/2))/86400, 12, 0.5); ok("beyond the Moon", rb > 384400, 1, 0)

# the thresholds
binf = lambda R: (np.sqrt(2) - 1)*(1 + 1/np.sqrt(R))
Rd = brentq(lambda x: binf(x) - g(x), 8, 15); ok("R dagger", Rd, 11.938765, 5e-6)
ok("limit formula", sum(bi(20, 1e14)), binf(20), 1e-6)
ok("third burn limit", (np.sqrt(2) - 1)/np.sqrt(9), np.sqrt(2/9) - 1/3, 1e-12)
Rr = [z.real for z in np.roots([1, -15, -9, -1]) if abs(z.imag) < 1e-9 and z.real > 0][0]; ok("R double dagger", Rr, 15.581719, 5e-7)
Rs, Ss = sp.symbols("R S", positive=True)
d2 = sp.sqrt(2/Ss)*(sp.sqrt(Rs/(Rs + Ss)) - sp.sqrt(1/(1 + Ss))); d3 = sp.sqrt(2*Ss/(Rs*(Rs + Ss))) - 1/sp.sqrt(Rs)
d1 = sp.sqrt(2*Ss/(1 + Ss)) - 1
ok("first term", sp.simplify(sp.diff(sp.sqrt(2/Ss)*sp.sqrt(Rs/(Rs + Ss)), Rs).subs(Ss, Rs) - 1/(4*Rs**sp.Rational(3, 2))), 0, 0)
ok("second term", sp.simplify(sp.diff(d3, Rs).subs(Ss, Rs) + 1/(4*Rs**sp.Rational(3, 2))), 0, 0)
C = d1 + d2 + d3; gs = sp.sqrt(2*Rs/(1 + Rs)) - 1 + (1 - sp.sqrt(2/(1 + Rs)))/sp.sqrt(Rs)
for v in (3.0, 15.0, 40.0):
    ok("dC/dS = g'(R) at S = R", (sp.diff(C, Ss).subs(Ss, Rs) - sp.diff(gs, Rs)).subs(Rs, v).evalf(), 0, 1e-10)
worst = min(g(x) - sum(bi(x, s)) for x in np.linspace(15.6, 100, 120) for s in np.geomspace(x*(1 + 1e-6), 1e4, 400))
ok("every r_b wins for 15.6 <= R <= 100", worst > 0, 1, 0)
worst_b = max(g(x) - sum(bi(x, s)) for x in np.linspace(2, 11.9, 60) for s in np.geomspace(x, 1e5, 300))
ok("Hohmann wins for R < 11.94", worst_b <= 1e-12, 1, 0)
ok("LEO to GEO", 42164/r1, 6.31, 5e-3)
# the 25 r_1 curve
Sf = 25.0
Rx = brentq(lambda x: sum(bi(x, Sf)) - g(x), 12, 16); ok("meets at", Rx, 14.10, 5e-3)
xs = np.linspace(14.2, 24.99, 2000); diffs = [g(x) - sum(bi(x, Sf)) for x in xs]; ok("largest gap", max(diffs), 0.0030, 5e-5)
i = int(np.argmax(diffs)); ok("near R", xs[i], 18.5, 0.15); ok("gap %", 100*diffs[i]/g(xs[i]), 0.56, 5e-3); ok("below throughout", min(diffs) > 0, 1, 0)

# the stopover
R, S = 30.0, 25.0
r2, rb = R*r1, S*r1
ok("r2", r2, 200344.11, 5e-3); ok("r_b", rb, 166953.43, 6e-3)
a1, a2 = (r1 + rb)/2, (rb + r2)/2
ok("a_t1", a1, 86815.78, 5e-3); ok("a_t2", a2, 183648.77, 5e-3)
rows = ((vc1, vis(r1, a1), 7.72576, 10.71370, 2.98794), (vis(rb, a1), vis(rb, a2), 0.42855, 1.61386, 1.18531), (vis(r2, a2), np.sqrt(mu/r2), 1.34488, 1.41052, 0.06564))
dv = []
for b, c, B, C_, D in rows:
    ok("before", b, B, 6e-6); ok("after", c, C_, 6e-6); ok("dv", c - b, D, 6e-6); dv.append(c - b)
tot = sum(dv); H = g(R)*vc1
ok("total", tot, 4.23889, 6e-6); ok("Hohmann", H, 4.07470, 6e-6); ok("more %", 100*(tot/H - 1), 4.03, 5e-3)
ts = (half(a1) + half(a2))/86400; th = half((r1 + r2)/2)/86400
ok("days", ts, 6.01, 5e-3); ok("Hohmann days", th, 1.92, 5e-3); ok("about three times", ts/th, 3, 0.2)
h1 = vis(r1, (r1 + r2)/2) - vc1; ok("Hohmann burn 1", h1, 3.02245, 6e-6); ok("saved m/s", 1e3*(h1 - dv[0]), 34.5, 0.05); ok("not climbed", r2 - rb, 33391, 0.5)
h2 = np.sqrt(mu/r2) - vis(r2, (r1 + r2)/2); ok("extra m/s", 1e3*(dv[1] + dv[2] - h2), 198.7, 0.05)
ok("speed ratio", vis(r1, a1)/vis(rb, a1), 25, 1e-9)
sto = lambda S_: (np.sqrt(2*S_/(1 + S_)) - 1) + (np.sqrt(2*R/(S_*(R + S_))) - np.sqrt(2/(S_*(1 + S_)))) + (1/np.sqrt(R) - np.sqrt(2*S_/(R*(R + S_))))
ok("stopover formula", sto(S)*vc1, tot, 1e-9); ok("ends: S = 1", sto(1.0), g(R), 1e-12); ok("ends: S = R", sto(R), g(R), 1e-12)
Sg = np.linspace(1.001, 29.999, 60000); vals = sto(Sg); i = int(np.argmax(vals))
ok("worst S", Sg[i], 4.02, 5e-3); ok("worst %", 100*(vals[i]/g(R) - 1), 34, 0.5); ok("always worse", vals.min() > g(R), 1, 0)
b40 = sum(bi(30, 40))*vc1; ok("r_b = 40 r_1", b40, 4.02458, 6e-6); ok("less %", 100*(1 - b40/H), 1.2, 0.05)

# the exercise
rm = 384400.0; am = (r1 + rm)/2; R = rm/r1
d1_, d2_ = vis(r1, am) - vc1, np.sqrt(mu/rm) - vis(rm, am)
ok("burn 1", d1_, 3.106, 5e-4); ok("burn 2", d2_, 0.830, 5e-4); ok("total", d1_ + d2_, 3.937, 6e-4); ok("days", half(am)/86400, 4.98, 5e-3); ok("R", R, 57.6, 0.05)
S = 2*R; ok("saving at 2 r_2, m/s", 1e3*(g(R) - sum(bi(R, S)))*vc1, 133, 0.5)
ok("time at 2 r_2, days", (half((r1 + 2*rm)/2) + half((rm + 2*rm)/2))/86400, 39, 0.5); ok("limit saving, m/s", 1e3*(g(R) - binf(R))*vc1, 315, 0.5)

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