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

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, r2 = 398600.4418, 6678.137, 42164.0
ve = 9.80665e-3*320                            # exhaust speed for I_sp = 320 s, km/s
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)

# the velocity triangle: the cost is least when the velocities are parallel
for ph in (0.0, 0.3, 1.0):
    ok("law of cosines", np.sqrt(9**2 + 7**2 - 2*9*7*np.cos(ph)), np.linalg.norm(9*np.array((np.cos(ph), np.sin(ph))) - (7, 0)), 1e-12)

# the Hohmann transfer
at = (r1 + r2)/2
vc1, vc2, vt1, vt2 = np.sqrt(mu/r1), np.sqrt(mu/r2), vis(r1, at), vis(r2, at)
ok("a_t", at, 24421.07, 5e-3); ok("e_t", (r2 - r1)/(r2 + r1), 0.72654, 5e-6)
ok("v_c1", vc1, 7.72576, 5e-6); ok("v_t1", vt1, 10.15149, 5e-6); ok("v_t2", vt2, 1.60784, 5e-6); ok("v_c2", vc2, 3.07467, 5e-6)
d1, d2 = vt1 - vc1, vc2 - vt2
ok("dv1", d1, 2.42573, 5e-6); ok("dv2", d2, 1.46682, 5e-6); ok("total", d1 + d2, 3.89255, 5e-6); ok("printed sum", 2.42573 + 1.46682, 3.89255, 1e-9)
ok("closed form dv1", np.sqrt(mu/r1)*(np.sqrt(2*r2/(r1 + r2)) - 1), d1, 1e-12)
ok("closed form dv2", np.sqrt(mu/r2)*(1 - np.sqrt(2*r1/(r1 + r2))), d2, 1e-12)
ok("t_H s", half(at), 18990, 0.5); ok("t_H h", half(at)/3600, 5.275, 5e-4)
ok("mass left, all", 100*left(d1 + d2), 28.9, 0.05); ok("mass left, second burn", 100*left(d2), 62.7, 0.05)
R_, r1_, r2_, mu_ = sp.symbols("R r1 r2 mu", positive=True)
ok("v_t1^2 identity", sp.simplify(mu_*(2/r1_ - 2/(r1_ + r2_)) - 2*mu_*r2_/(r1_*(r1_ + r2_))), 0, 0)
ok("v_t2^2 identity", sp.simplify(mu_*(2/r2_ - 2/(r1_ + r2_)) - 2*mu_*r1_/(r2_*(r1_ + r2_))), 0, 0)

# the cost against R
g = lambda R: np.sqrt(2*R/(1 + R)) - 1 + (1 - np.sqrt(2/(1 + R)))/np.sqrt(R)
ok("R of LEO-GEO", r2/r1, 6.31, 5e-3); ok("g = total/v_c1", g(r2/r1), (d1 + d2)/vc1, 1e-12)
for R, v in ((6.31, 0.5038), (15, 0.5362), (100, 0.4931), (1e12, 0.4142)):
    ok(f"g({R})", g(R), v, 6e-5)
gs = sp.sqrt(2*R_/(1 + R_)) - 1 + (1 - sp.sqrt(2/(1 + R_)))/sp.sqrt(R_)
dg_note = 1/(sp.sqrt(2*R_)*(1 + R_)**sp.Rational(3, 2)) - 1/(2*R_**sp.Rational(3, 2)) + (1 + 2*R_)/(sp.sqrt(2)*R_**sp.Rational(3, 2)*(1 + R_)**sp.Rational(3, 2))
for Rv in (2.0, 7.0, 30.0):
    ok("dg/dR as printed", (sp.diff(gs, R_) - dg_note).subs(R_, Rv).evalf(), 0, 1e-12)
ok("the cubic", sp.expand(2*(1 + 3*R_)**2 - (1 + R_)**3 + (R_**3 - 15*R_**2 - 9*R_ - 1)), 0, 0)
roots = np.roots([1, -15, -9, -1]); pos = [z.real for z in roots if abs(z.imag) < 1e-9 and z.real > 0]
ok("one positive root", len(pos), 1, 0); ok("the root", pos[0], 15.581719, 5e-7)
m = minimize_scalar(lambda R: -g(R), bounds=(5, 40), method="bounded", options=dict(xatol=1e-8))
ok("numerical maximum at", m.x, 15.581719, 2e-4); ok("maximum", g(pos[0]), 0.536258, 5e-7); ok("figure label", g(pos[0]), 0.5363, 5e-5)
ok("the bound from 300 km", g(pos[0])*vc1, 4.143, 5e-4)
ok("escape burn in v_c1", (np.sqrt(2*mu/r1) - vc1)/vc1, np.sqrt(2) - 1, 1e-12)
ok("time to the Moon's distance, days", half((r1 + 384400)/2)/86400, 5, 0.03)

# coaxial, aligned and anti-aligned
a1, e1, a2, e2 = 12000.0, 0.30, 42000.0, 0.50


def route(rA, rB, oA, oB):
    a = (rA + rB)/2
    return a, abs(vis(rA, a) - vis(rA, oA)), abs(vis(rB, a) - vis(rB, oB)), half(a)/3600


rp1, ra1, rp2, ra2 = a1*(1 - e1), a1*(1 + e1), a2*(1 - e2), a2*(1 + e2)
for n, x, y in (("rp1", rp1, 8400), ("ra1", ra1, 15600), ("rp2", rp2, 21000), ("ra2", ra2, 63000)):
    ok(n, x, y, 1e-9)
tab = (("A1", rp1, ra2, 35700, 1.29675, 0.55850, 1.85525, 9.324), ("A2", ra1, rp2, 18300, 1.18573, 1.31337, 2.49909, 3.422),
       ("T2", rp1, rp2, 14700, 0.37924, 2.04250, 2.42174, 2.464), ("T1", ra1, ra2, 39300, 2.17083, 0.19386, 2.36469, 10.769))
res = {}
for n, ra_, rb_, A, D1, D2, T, H in tab:
    a, x, y, h = route(ra_, rb_, a1, a2)
    res[n] = (x + y, h)
    ok(n + " a_t", a, A, 1e-9); ok(n + " dv1", x, D1, 6e-6); ok(n + " dv2", y, D2, 6e-6); ok(n + " total", x + y, T, 6e-6); ok(n + " h", h, H, 6e-4)
ok("A1 cheaper, m/s", 1e3*(res["A2"][0] - res["A1"][0]), 644, 0.5); ok("A2 faster", res["A1"][1]/res["A2"][1], 2.72, 5e-3)
ok("mass left A1", 100*left(res["A1"][0]), 55.4, 0.05); ok("mass left A2", 100*left(res["A2"][0]), 45.1, 0.05)
ok("A1 more mass, percent", 100*(left(res["A1"][0])/left(res["A2"][0]) - 1), 23, 0.5); ok("A2 sooner, h", res["A1"][1] - res["A2"][1], 5.9, 0.05)
ok("T1 cheaper, m/s", 1e3*(res["T2"][0] - res["T1"][0]), 57, 0.5); ok("T1 cheaper, percent", 100*(res["T2"][0] - res["T1"][0])/res["T2"][0], 2.4, 0.05)
ok("T2 faster", res["T1"][1]/res["T2"][1], 4.37, 5e-3); ok("orientation, percent", 100*(res["T1"][0]/res["A1"][0] - 1), 27, 0.5); ok("T1 longer, h", res["T1"][1] - res["T2"][1], 8.3, 0.05)
c1, c2 = a1*e1, a2*e2
ok("c1 + c2", c1 + c2, 24600, 1e-9); ok("= difference of a_t", 39300 - 14700, c1 + c2, 1e-9); ok("c2 - c1", 35700 - 18300, c2 - c1, 1e-9)
# e1 = 0.45
rp, ra = a1*0.55, a1*1.45
ok("rp", rp, 6600, 1e-9); ok("ra", ra, 17400, 1e-9)
aT2, x, y, h = route(rp, rp2, a1, a2); ok("T2 a_t", aT2, 13800, 1e-9); ok("T2 dv1", x, 0.22870, 6e-6); ok("T2 dv2", y, 2.32292, 6e-6); ok("T2", x + y, 2.55162, 6e-6); ok("T2 h", h, 2.241, 6e-4)
aT1, x, y, h = route(ra, ra2, a1, a2); ok("T1 a_t", aT1, 40200, 1e-9); ok("T1 dv1", x, 2.44215, 6e-6); ok("T1 dv2", y, 0.12377, 6e-6); ok("T1", x + y, 2.56592, 6e-6); ok("T1 h", h, 11.141, 6e-4)

# the proof: the cyclic sum against a direct evaluation, and A1 cheaper on random orbits
Phi = lambda s: np.sqrt(2*mu/s)
L = lambda s, t: (s - t)*Phi(s + t)
C3 = lambda s, t, k: L(s, t) + L(t, k) + L(k, s)
rng = np.random.default_rng(5)
worst = 1e9; bad = 0
for _ in range(4000):
    p1 = rng.uniform(7000, 20000); q1 = rng.uniform(p1, 4*p1); p2 = rng.uniform(1.01*p1, 6*p1); q2 = rng.uniform(max(p2, 1.01*q1), 8*p2)
    aa1, aa2 = (p1 + q1)/2, (p2 + q2)/2
    A1 = route(p1, q2, aa1, aa2); A2 = route(q1, p2, aa1, aa2)
    x1, y1, x2, y2 = 1/p1, 1/q1, 1/p2, 1/q2
    cyc = L(x1, y1) + L(y1, x2) + L(x2, y2) + L(y2, x1)
    bad += abs((A2[1] + A2[2]) - (A1[1] + A1[2]) - cyc) > 1e-9
    bad += abs(cyc - (C3(x1, y1, y2) + C3(y1, x2, y2))) > 1e-9
    worst = min(worst, cyc)
ok("cyclic sum = dv_A2 - dv_A1 = C3 + C3 (4000 cases)", bad, 0, 0); ok("A1 cheaper in all of them", worst > 0, 1, 0)
# anti-aligned: the sign when one orbit is a circle


def anti(r1a, r1p, r2a, r2p):
    A1_, A2_ = (r1a + r1p)/2, (r2a + r2p)/2
    aT1_, aT2_ = (r1a + r2a)/2, (r1p + r2p)/2
    return (vis(r1a, aT1_) - vis(r1a, A1_) + vis(r2a, A2_) - vis(r2a, aT1_)) - (vis(r1p, aT2_) - vis(r1p, A1_) + vis(r2p, A2_) - vis(r2p, aT2_))


s1 = s2 = True
for _ in range(2000):
    rc = rng.uniform(7000, 15000); q = rng.uniform(rc*1.05, 4*rc); w = rng.uniform(q, 10*q); s1 &= anti(rc, rc, w, q) < 0
    rC = rng.uniform(30000, 60000); q = rng.uniform(7000, rC/1.05); w = rng.uniform(q, rC/1.01); s2 &= anti(w, q, rC, rC) > 0
ok("orbit 1 a circle: T1 cheaper (2000 cases)", s1, 1, 0); ok("orbit 2 a circle: T2 cheaper (2000 cases)", s2, 1, 0)

# crossing orbits
Rp1, Ra1, Rp2, Ra2 = 8000.0, 30000.0, 12000.0, 24000.0
A1_, E1, A2_, E2 = (Rp1 + Ra1)/2, (Ra1 - Rp1)/(Ra1 + Rp1), (Rp2 + Ra2)/2, (Ra2 - Rp2)/(Ra2 + Rp2)
P1, P2 = A1_*(1 - E1**2), A2_*(1 - E2**2)
ok("a1", A1_, 19000, 1e-9); ok("e1", E1, 0.578947, 5e-7); ok("p1", P1, 12631.58, 5e-3); ok("a2", A2_, 18000, 1e-9); ok("e2", E2, 1/3, 1e-12); ok("p2", P2, 16000, 1e-9)
cf = (P2 - P1)/(P1*E2 - P2*E1)
ok("numerator", P2 - P1, 3368.42, 5e-3); ok("denominator", P1*E2 - P2*E1, -5052.63, 5e-3); ok("cos f*", cf, -0.666667, 5e-7)
fs = np.arccos(cf); ok("f*", np.degrees(fs), 131.81, 5e-3); rs = P1/(1 + E1*cf); ok("r*", rs, 20571.43, 5e-3); ok("same on orbit 2", P2/(1 + E2*cf), rs, 1e-8)
h1, h2 = np.sqrt(mu*P1), np.sqrt(mu*P2)
vr1, vt1_, vr2, vt2_ = mu/h1*E1*np.sin(fs), h1/rs, mu/h2*E2*np.sin(fs), h2/rs
for n, x, y in (("vr1", vr1, 2.42406), ("vth1", vt1_, 3.44932), ("v1", np.hypot(vr1, vt1_), 4.21590), ("vr2", vr2, 1.24008), ("vth2", vt2_, 3.88208), ("v2", np.hypot(vr2, vt2_), 4.07533)):
    ok(n, x, y, 6e-6)
g1, g2 = np.degrees(np.arctan2(vr1, vt1_)), np.degrees(np.arctan2(vr2, vt2_))
ok("gamma1", g1, 35.098, 5e-4); ok("gamma2", g2, 17.715, 5e-4); ok("turn", g1 - g2, 17.38, 5e-3)
single = np.hypot(vr2 - vr1, vt2_ - vt1_); ok("single impulse", single, 1.26058, 6e-5)
a, x, y, h = route(Rp1, Ra2, A1_, A2_); ok("route 1 a", a, 16000, 1e-9); ok("dv1", x, 0.22459, 6e-6); ok("dv2", y, 0.44580, 6e-6); ok("total", x + y, 0.67039, 6e-6); ok("h", h, 2.797, 6e-4)
a, x, y, h = route(Ra1, Rp2, A1_, A2_); ok("route 2 a", a, 21000, 1e-9); ok("dv1", x, 0.39018, 6e-6); ok("dv2", y, 0.23358, 6e-6); ok("total", x + y, 0.62376, 6e-6); ok("h", h, 4.206, 6e-4)
cheap = x + y; ok("ratio", single/cheap, 2.02, 5e-3)
vb = (np.hypot(vr1, vt1_) + np.hypot(vr2, vt2_))/2; ok("mean speed", vb, 4.146, 5e-4)
turn = 2*vb*np.sin(np.radians(g1 - g2)/2); ok("pure turn", turn, 1.253, 5e-4); ok("fraction", 100*turn/single, 99, 0.6)
T_1 = 2*half(A1_); ok("period of orbit 1, h", T_1/3600, 7.240, 5e-4)
Ec = 2*np.arctan(np.sqrt((1 - E1)/(1 + E1))*np.tan(fs/2)); t_ac = T_1/2 - (Ec - E1*np.sin(Ec))/(2*np.pi)*T_1
ok("crossing to apoapsis, h", t_ac/3600, 2.305, 5e-4); coast = half(21000.0); ok("coast, h", coast/3600, 4.206, 5e-4)
ok("saved, case 1", (coast - t_ac)/3600, 1.90, 5e-3); ok("saved, case 2", (t_ac + coast)/3600, 6.51, 5e-3); ok("saved, case 3", (T_1 - t_ac + coast)/3600, 9.14, 5e-3)
# the three cases, by brute force over the starting point on orbit 1


def tof(f0, f1):
    M = lambda f: (lambda E_: E_ - E1*np.sin(E_))(2*np.arctan2(np.sqrt(1 - E1)*np.sin(f/2), np.sqrt(1 + E1)*np.cos(f/2)))
    return ((M(f1) - M(f0)) % (2*np.pi))/(2*np.pi)*T_1


for f0, want in ((np.radians(150), 1.90), (np.radians(0), 6.51), (np.radians(200), 9.14)):
    ok("brute force", (tof(f0, np.pi) + coast - min(tof(f0, fs), tof(f0, -fs)))/3600, want, 5e-3)
ok("mass left, single", 100*left(single), 66.9, 0.05); ok("mass left, two burns", 100*left(cheap), 82.0, 0.05)

# the optimality proof
for (r_a, r_b) in ((r1, r2),):
    hH = np.sqrt(2*mu*r_a*r_b/(r_a + r_b))
    ok("h_H is the ellipse's h", hH, vt1*r1, 1e-8)
    eps = -mu/(r_a + r_b)
    dv = lambda r, e_, h_: np.sqrt(2*e_ + 3*mu/r - 2*h_*np.sqrt(mu)/r**1.5)
    ok("eq-dv-eh gives dv1", dv(r1, eps, hH), d1, 1e-8); ok("eq-dv-eh gives dv2", dv(r2, eps, hH), d2, 1e-8)
    F = lambda r, h_: h_*h_/(2*r*r) - mu/r
    ok("F(r1) = F(r2) at the corner", F(r1, hH), F(r2, hH), 1e-9); ok("F = energy", F(r1, hH), eps, 1e-9)
    hs = np.linspace(0.6*hH, 1.4*hH, 801); cost = [dv(r1, max(F(r1, h_), F(r2, h_)), h_) + dv(r2, max(F(r1, h_), F(r2, h_)), h_) for h_ in hs]
    ok("minimum at the corner", hs[int(np.argmin(cost))]/hH, 1, 1e-3); ok("minimum", min(cost), 3.8926, 5e-4)
# exact identity, by brute force: dv from the velocity components
for _ in range(50):
    r_ = rng.uniform(6000, 50000); vr_ = rng.uniform(-3, 3); vp = rng.uniform(1, 9)
    eps_ = (vr_**2 + vp**2)/2 - mu/r_; h_ = r_*vp
    if abs(np.sqrt(2*eps_ + 3*mu/r_ - 2*h_*np.sqrt(mu)/r_**1.5) - np.hypot(vr_, vp - np.sqrt(mu/r_))) > 1e-8:
        fails += 1; print("FAIL eq-dv-eh")

# exercises
rP = 6378.137 + 100; at_ = (rP + r2)/2
ok("going down", vc2 - vis(r2, at_), 1.488, 5e-4)
ag = (r2 + r2 + 300)/2
gy = (vis(r2, ag) - vc2) + (np.sqrt(mu/(r2 + 300)) - vis(r2 + 300, ag))
ok("graveyard, m/s", 1e3*gy, 10.9, 0.05); ok("ratio", (vc2 - vis(r2, at_))/gy, 137, 0.5)

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