"""Checks every result quoted in orbital-maneuvers-analytic.qmd.

Each problem gets (i) a symbolic check of the closed form with sympy, where one exists,
and (ii) a brute-force numerical check of the inequality or optimum over a grid.
Run:  python codes/maneuvers-analytic-verify.py      (prints PASS/FAIL per claim)
"""
import numpy as np
import sympy as sp
from scipy.optimize import brentq, minimize_scalar

RESULTS = []


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


mu = 398600.4418  # km^3/s^2, only for numerical spot checks

# ---------------------------------------------------------------- 1. Hohmann: which burn is larger
r1, r2, m = sp.symbols("r1 r2 mu", positive=True)
dv1 = sp.sqrt(m / r1) * (sp.sqrt(2 * r2 / (r1 + r2)) - 1)
dv2 = sp.sqrt(m / r2) * (1 - sp.sqrt(2 * r1 / (r1 + r2)))
claim = sp.sqrt(m) * (sp.sqrt(2 * (r1 + r2)) - sp.sqrt(r1) - sp.sqrt(r2)) / sp.sqrt(r1 * r2)
diff = sp.simplify(sp.expand(dv1 - dv2 - claim))
num_ok = all(abs(float((dv1 - dv2 - claim).subs({m: mu, r1: 7000.0, r2: x}))) < 1e-12
             for x in [7000.0, 8000.0, 42164.0, 1e6])
check("1  dv1-dv2 = sqrt(mu)[sqrt(2(r1+r2)) - sqrt r1 - sqrt r2]/sqrt(r1 r2)", diff == 0 or num_ok)
R = np.linspace(1.0001, 1e4, 200000)
d1 = np.sqrt(2 * R / (1 + R)) - 1
d2 = (1 - np.sqrt(2 / (1 + R))) / np.sqrt(R)
check("1  dv1 > dv2 for all R>1 (grid)", np.all(d1 > d2))
check("1  dv1/dv2 -> 1 as R->1", abs(d1[0] / d2[0] - 1) < 1e-3)
# inward transfer: the burn at the inner circle is the same magnitude (time reversal)

# ---------------------------------------------------------------- 2. Hohmann maximum
Rs = sp.symbols("R", positive=True)
g = sp.sqrt(2 * Rs / (1 + Rs)) - 1 + 1 / sp.sqrt(Rs) - sp.sqrt(2 / (Rs * (1 + Rs)))
gp = sp.diff(g, Rs)
# claim: 2 R^{3/2}(1+R)^{3/2} g'(R) = sqrt2 (1+3R) - (1+R)^{3/2}
lhs = sp.simplify(2 * Rs**sp.Rational(3, 2) * (1 + Rs)**sp.Rational(3, 2) * gp)
rhs = sp.sqrt(2) * (1 + 3 * Rs) - (1 + Rs)**sp.Rational(3, 2)
check("2  2R^{3/2}(1+R)^{3/2} g' = sqrt2(1+3R)-(1+R)^{3/2}",
      all(abs(float((lhs - rhs).subs(Rs, x))) < 1e-10 for x in [1.5, 3.0, 15.0, 200.0]))
cubic = sp.Poly(sp.expand((1 + Rs)**3 - 2 * (1 + 3 * Rs)**2), Rs)
check("2  (1+R)^3 - 2(1+3R)^2 = R^3 - 15R^2 - 9R - 1", cubic.as_expr() == Rs**3 - 15 * Rs**2 - 9 * Rs - 1)
roots = np.roots([1, -15, -9, -1])
real = sorted(r.real for r in roots if abs(r.imag) < 1e-12)
check("2  cubic has exactly one root > 1, R=15.5817", len([x for x in real if x > 1]) == 1
      and abs(max(real) - 15.581719) < 1e-5)
Rstar = max(real)
gstar = float(g.subs(Rs, Rstar))
check("2  max g = 0.536258", abs(gstar - 0.536258) < 1e-6)
# arrival burn alone
h_arr = 1 / sp.sqrt(Rs) - sp.sqrt(2 / (Rs * (1 + Rs)))
hp = sp.diff(h_arr, Rs)
lhs = sp.simplify(2 * Rs**sp.Rational(3, 2) * (1 + Rs)**sp.Rational(3, 2) * hp)
rhs = sp.sqrt(2) * (1 + 2 * Rs) - (1 + Rs)**sp.Rational(3, 2)
check("2  arrival: 2R^{3/2}(1+R)^{3/2} h' = sqrt2(1+2R)-(1+R)^{3/2}",
      all(abs(float((lhs - rhs).subs(Rs, x))) < 1e-10 for x in [1.5, 3.0, 15.0]))
cub2 = sp.expand((1 + Rs)**3 - 2 * (1 + 2 * Rs)**2)
check("2  arrival cubic R^3 - 5R^2 - 5R - 1", cub2 == Rs**3 - 5 * Rs**2 - 5 * Rs - 1)
r2roots = [r.real for r in np.roots([1, -5, -5, -1]) if abs(r.imag) < 1e-12]
Rarr = max(r2roots)
check("2  arrival cubic has one root > 1, R = 5.8794", len([x for x in r2roots if x > 1]) == 1
      and abs(Rarr - 5.879362) < 1e-6)
check("2  arrival cubic has no rational root", sp.Poly(cub2, Rs).factor_list()[1][0][0].degree() == 3)


# ---------------------------------------------------------------- 3. aligned coaxial ellipses: which route
def H(x, y):  # angular momentum of the conic with apsides x, y
    return np.sqrt(2 * mu * x * y / (x + y))


def Phi(s):
    return np.sqrt(2 * mu / s)


def L(x, y):
    return (x - y) * Phi(x + y)


rng = np.random.default_rng(1)
ok_cyc = ok_sign = True
n_cross = 0
for _ in range(200000):
    rp1, ra1, rp2, ra2 = rng.uniform(6500, 4e5, 4)
    rp1, ra1 = min(rp1, ra1), max(rp1, ra1)
    rp2, ra2 = min(rp2, ra2), max(rp2, ra2)
    if not (rp1 < rp2 and ra1 < ra2):
        continue
    n_cross += ra1 > rp2
    h1, h2 = H(rp1, ra1), H(rp2, ra2)
    A1 = (H(rp1, ra2) - h1) / rp1 + (h2 - H(rp1, ra2)) / ra2      # p1 -> a2
    A2 = (H(ra1, rp2) - h1) / ra1 + (h2 - H(ra1, rp2)) / rp2      # a1 -> p2
    a, b, c, d = 1 / rp1, 1 / ra1, 1 / rp2, 1 / ra2
    cyc = L(a, b) + L(b, c) + L(c, d) + L(d, a)
    ok_cyc &= abs((A2 - A1) - cyc) < 1e-9
    ok_sign &= A2 > A1
check("3  A2 - A1 = cyclic sum L(a,b)+L(b,c)+L(c,d)+L(d,a)", ok_cyc)
check(f"3  A1 cheaper whenever rp1<rp2, ra1<ra2 ({n_cross} samples with ra1>rp2)", ok_sign)
# anti-aligned: T1 - T2 is the cycle (a,b,d,c); its sign flips between the two circular limits
def anti(rp1, ra1, rp2, ra2):
    h1, h2 = H(rp1, ra1), H(rp2, ra2)
    T2 = abs(H(rp1, rp2) - h1) / rp1 + abs(h2 - H(rp1, rp2)) / rp2    # peri -> peri
    T1 = abs(H(ra1, ra2) - h1) / ra1 + abs(h2 - H(ra1, ra2)) / ra2    # apo -> apo
    a, b, c, d = 1 / rp1, 1 / ra1, 1 / rp2, 1 / ra2
    return T1 - T2, L(a, b) + L(b, d) + L(d, c) + L(c, a)
dA, cA = anti(8000, 8000, 21000, 63000)      # inner circular
dB, cB = anti(8000, 15600, 30000, 30000)     # outer circular
check("3  anti-aligned: T1-T2 = cycle (a,b,d,c)", abs(dA - cA) < 1e-9 and abs(dB - cB) < 1e-9)
check("3  anti-aligned: T1 cheaper if inner circular, T2 cheaper if outer circular", dA < 0 < dB)
d2, _ = anti(8400, 15600, 21000, 63000)
print("     deck Example 3: T1-T2 =", d2, "(T1 cheaper)")


# ---------------------------------------------------------------- 4. one-burn circularisation
e_, x_ = sp.symbols("e x", positive=True)
F = 3 * x_**2 - 2 * x_**3 - (1 - e_**2)
s_, t_ = sp.sqrt(1 - e_), sp.sqrt(1 + e_)
check("4  F(sqrt(1-e)) = (1-e)(1-sqrt(1-e))^2",
      sp.simplify(F.subs(x_, s_) - (1 - e_) * (1 - s_)**2) == 0)
check("4  F(sqrt(1+e)) = (1+e)(sqrt(1+e)-1)^2",
      sp.simplify(F.subs(x_, t_) - (1 + e_) * (t_ - 1)**2) == 0)
check("4  F(1) = e^2", sp.simplify(F.subs(x_, 1) - e_**2) == 0)
# direct numerical check of the one-burn cost against the formula
ok = True
for a, e in [(10000, 0.3), (20000, 0.7), (8000, 0.05)]:
    p = a * (1 - e * e)
    for f in np.linspace(0, np.pi, 50):
        r = p / (1 + e * np.cos(f))
        vr = np.sqrt(mu / p) * e * np.sin(f)
        vt = np.sqrt(mu / p) * (1 + e * np.cos(f))
        dv = np.hypot(vr, vt - np.sqrt(mu / r))
        xx = np.sqrt(1 + e * np.cos(f))
        ok &= abs(dv - np.sqrt(mu / p) * np.sqrt(3 * xx**2 - 2 * xx**3 - (1 - e * e))) < 1e-10
check("4  dv(r)^2 = (mu/p)(3x^2-2x^3-(1-e^2)), x^2 = p/r... (here x^2 = 1+e cos f)", ok)
ee = np.linspace(1e-4, 0.9999, 10000)
check("4  apoapsis cheaper than periapsis for all e",
      np.all(np.sqrt(1 - ee) * (1 - np.sqrt(1 - ee)) < np.sqrt(1 + ee) * (np.sqrt(1 + ee) - 1)))

# ---------------------------------------------------------------- 5. apse-line rotation crossover
def eta_c(e):
    s = np.sqrt(1 - e)
    return 2 * np.degrees(np.arcsin(s / (1 + s)))


def eta_c_direct(a, e):
    p = a * (1 - e * e)
    ra = a * (1 + e)
    alt = 2 * (np.sqrt(mu / ra) - np.sqrt(mu * (2 / ra - 1 / a)))
    return 2 * np.degrees(np.arcsin(alt / (2 * np.sqrt(mu / p) * e)))


check("5  eta_c(e=0.4) = 51.8 deg (deck)", abs(eta_c(0.4) - 51.8) < 0.05)
check("5  closed form = direct crossover", all(abs(eta_c(e) - eta_c_direct(12000, e)) < 1e-9
                                              for e in [0.01, 0.3, 0.4, 0.9]))
check("5  eta_c -> 60 deg as e -> 0", abs(eta_c(1e-9) - 60) < 1e-6)
check("5  eta_c monotone decreasing", np.all(np.diff(eta_c(ee)) < 0))

# ---------------------------------------------------------------- 6. phasing
X = sp.symbols("x", positive=True)
lower = 1 - sp.sqrt(2 - (1 - X)**sp.Rational(-2, 3))       # each burn / v_c, chaser behind
upper = sp.sqrt(2 - (1 + X)**sp.Rational(-2, 3)) - 1       # chaser ahead
sl = sp.series(lower, X, 0, 4).removeO()
su = sp.series(upper, X, 0, 4).removeO()
print("     behind:", sp.expand(sl), "   ahead:", sp.expand(su))
check("6  behind: x/3 + x^2/3 + ...", sp.expand(sl).coeff(X, 1) == sp.Rational(1, 3)
      and sp.expand(sl).coeff(X, 2) == sp.Rational(1, 3))
check("6  ahead:  x/3 - x^2/3 + ...", sp.expand(su).coeff(X, 1) == sp.Rational(1, 3)
      and sp.expand(su).coeff(X, 2) == -sp.Rational(1, 3))
# product dv_total * t / (r dtheta) with x = dtheta/(2 pi k):  dv_total = 2 v_c * lower, t = k T (1-x)
# = 2 v_c * lower * k (2 pi r / v_c)(1-x) / (r * 2 pi k x) = 2 lower (1-x)/x
prod = sp.series(2 * lower * (1 - X) / X, X, 0, 3).removeO()
print("     dv*t/(r dtheta) =", sp.expand(prod))
check("6  dv*t = (2/3) r dtheta (1 + 0*x + O(x^2))", sp.expand(prod).coeff(X, 0) == sp.Rational(2, 3)
      and sp.expand(prod).coeff(X, 1) == 0)
# Example 6 of the deck
r, dth = 6778.0, np.radians(10)
T = 2 * np.pi * np.sqrt(r**3 / mu)
aph = r * (1 - dth / (2 * np.pi))**(2 / 3)
dv = 2 * (np.sqrt(mu / r) - np.sqrt(mu * (2 / r - 1 / aph)))
check("6  Example 6: dv*t vs (2/3) r dtheta within 0.05%",
      abs(dv * T * (1 - dth / (2 * np.pi)) / (2 / 3 * r * dth) - 1) < 5e-4)
print("     Example 6:", dv * T * (1 - dth / (2 * np.pi)), (2 / 3) * r * dth)
# minimum k for a perigee floor
rmin = 6378.137 + 200
xmax = 1 - ((r + rmin) / (2 * r))**1.5
kmin = dth / (2 * np.pi * xmax)
a_k = r * (1 - dth / (2 * np.pi * kmin))**(2 / 3)
check("6  k_min formula puts perigee exactly on the floor", abs(2 * a_k - r - rmin) < 1e-8)

# ---------------------------------------------------------------- 7. plane change via a high apoapsis
K, S = sp.symbols("K s", positive=True)
C = 2 * (sp.sqrt(2 * K / (1 + K)) - 1) + 2 * S * sp.sqrt(2 / (K * (1 + K)))
dC = sp.diff(C, K)
claim = 2 / (sp.sqrt(2) * K**sp.Rational(3, 2) * (1 + K)**sp.Rational(3, 2)) * (K - S * (1 + 2 * K))
check("7  dC/dK = sqrt2 [K - s(1+2K)] / (K(1+K))^{3/2}",
      all(abs(float((dC - claim).subs({K: k, S: s}))) < 1e-12 for k in [1.0, 2.5, 30.0] for s in [0.2, 0.4, 0.7]))
Kopt = S / (1 - 2 * S)
Copt = sp.simplify(C.subs(K, Kopt))
check("7  C(K*) = 4 sqrt(2 s(1-s)) - 2",
      all(abs(float((Copt - (4 * sp.sqrt(2 * S * (1 - S)) - 2)).subs(S, s))) < 1e-12 for s in [0.34, 0.4, 0.45, 0.49]))
Cn = sp.lambdify((K, S), C, "numpy")
ok = True
for th in np.radians(np.linspace(1, 179, 357)):
    s = np.sin(th / 2)
    Kg = np.logspace(0, 6, 20001)
    best = Kg[np.argmin(Cn(Kg, s))]
    if s <= 1 / 3:
        ok &= best == 1.0
    elif s < 0.5 - 1e-9:
        ok &= abs(best / (s / (1 - 2 * s)) - 1) < 2e-3
    else:
        ok &= best == Kg[-1]
check("7  optimum: K=1 (s<=1/3), K=s/(1-2s) (1/3<s<1/2), K->inf (s>=1/2)", ok)
check("7  thresholds 38.94, 48.94, 60 deg",
      abs(2 * np.degrees(np.arcsin(1 / 3)) - 38.942) < 1e-3
      and abs(2 * np.degrees(np.arcsin(np.sqrt(2) - 1)) - 48.94) < 1e-2)

# ---------------------------------------------------------------- 8. folding the plane change in
vc1, vt1, vt2, vc2, d, D = sp.symbols("v_c1 v_t1 v_t2 v_c2 delta Di", positive=True)
dv2exp = sp.sqrt(vt2**2 + vc2**2 - 2 * vt2 * vc2 * sp.cos(D))
ser = sp.series(dv2exp.subs({vt2: 1.6, vc2: 3.1}), D, 0, 3).removeO()
check("8  dv2 = (vc2-vt2) + vt2 vc2 Di^2 / (2(vc2-vt2)) + O(Di^4)",
      abs(float(ser.coeff(D, 2)) - 1.6 * 3.1 / (2 * 1.5)) < 1e-12 and abs(float(ser.coeff(D, 0)) - 1.5) < 1e-12)


def split_opt(Rr, Di):
    v_c1, v_c2 = 1.0, 1 / np.sqrt(Rr)
    v_t1 = np.sqrt(2 * Rr / (1 + Rr))
    v_t2 = v_t1 / Rr
    d1 = lambda dd: np.sqrt(v_c1**2 + v_t1**2 - 2 * v_c1 * v_t1 * np.cos(dd))
    d2 = lambda dd: np.sqrt(v_t2**2 + v_c2**2 - 2 * v_t2 * v_c2 * np.cos(Di - dd))
    der = lambda dd: v_c1 * v_t1 * np.sin(dd) / d1(dd) - v_t2 * v_c2 * np.sin(Di - dd) / d2(dd)
    x = brentq(der, 1e-15 * Di, Di * (1 - 1e-12), xtol=1e-18)
    D1, D2 = v_t1 - v_c1, v_c2 - v_t2
    return x / Di, 1 / (1 + Rr**1.5 * D2 / D1)


ok = True
for Rr in [1.05, 2.0, 6.3136, 20.0]:
    got, pred = split_opt(Rr, np.radians(0.01))
    ok &= abs(got / pred - 1) < 1e-5
check("8  small-angle split delta/Di = 1/(1+R^{3/2} dv2/dv1)", ok)
print("     LEO-GEO fraction (small angle, 28.5 deg):", split_opt(42164 / 6678.137, np.radians(0.01)),
      split_opt(42164 / 6678.137, np.radians(28.5)))

# ---------------------------------------------------------------- 9. node change
def hhat(i, O):
    return np.array([np.sin(i) * np.sin(O), -np.sin(i) * np.cos(O), np.cos(i)])


ok1 = ok2 = True
for i in np.radians([5, 30, 51.6, 89, 120, 170]):
    for dO in np.radians([1, 10, 90, 170]):
        h1, h2 = hhat(i, -dO / 2), hhat(i, dO / 2)
        alpha = np.arccos(np.clip(h1 @ h2, -1, 1))
        ok1 &= abs(np.sin(alpha / 2) - np.sin(i) * np.sin(dO / 2)) < 1e-12
        x = np.cross(h1, h2)
        x /= np.linalg.norm(x)          # northern intersection (x_K >= 0 for 0<i<pi)
        n1 = np.array([np.cos(-dO / 2), np.sin(-dO / 2), 0])
        n2 = np.array([np.cos(dO / 2), np.sin(dO / 2), 0])
        u1 = np.arctan2(x @ np.cross(h1, n1), x @ n1)   # angle from n1 to x about h1
        u2 = np.arctan2(x @ np.cross(h2, n2), x @ n2)
        ok2 &= abs(u1 - (np.pi / 2 + np.arctan(np.cos(i) * np.tan(dO / 2)))) < 1e-10
        ok2 &= abs(u2 - (np.pi / 2 - np.arctan(np.cos(i) * np.tan(dO / 2)))) < 1e-10
check("9  sin(alpha/2) = sin i sin(dOmega/2)", ok1)
check("9  burn at u1 = 90+arctan(cos i tan(dO/2)), u2 = 90-arctan(...)", ok2)

# ---------------------------------------------------------------- 10. Oberth dive before escape
def dive(rho, V):
    return (1 - np.sqrt(2 * rho / (1 + rho))) + np.sqrt(V**2 + 2 / rho) - np.sqrt(2 / (rho * (1 + rho)))


rho = np.logspace(-6, 0, 200001)
ok = True
for V in np.linspace(0.0, 4.0, 401):
    c = dive(rho, V)
    direct = np.sqrt(V**2 + 2) - 1
    if V < np.sqrt(2) - 1e-9:
        ok &= c.min() >= direct - 1e-12          # never better than not diving
    elif V > np.sqrt(2) + 1e-9:
        ok &= np.all(np.diff(c) > -1e-15)        # cost rises with rho: deeper is cheaper
check("10 dive pays iff v_inf > sqrt2 v_c; then deeper is always cheaper", ok)
rr, VV = sp.symbols("rho V", positive=True)
Cd = (1 - sp.sqrt(2 * rr / (1 + rr))) + sp.sqrt(VV**2 + 2 / rr) - sp.sqrt(2 / (rr * (1 + rr)))
slope1 = sp.simplify(sp.diff(Cd, rr).subs(rr, 1))
check("10 dC/drho at rho=1 = 1/2 - 1/sqrt(V^2+2)",
      sp.simplify(slope1 - (sp.Rational(1, 2) - 1 / sp.sqrt(VV**2 + 2))) == 0)
check("10 limit rho->0 cost = v_c", abs(dive(1e-12, 2.0) - 1) < 1e-5)

# ---------------------------------------------------------------- steps added in the pedagogical rewrite
# 1: second proof via energy, dv1/dv2 = R (vt2+vc2)/(vc1+vt1), and vc1+vt1 < R(vc2+vt2)
Rg = np.linspace(1.001, 1e3, 20000)
vc1, vc2 = 1.0, 1 / np.sqrt(Rg)
vt1, vt2 = np.sqrt(2 * Rg / (1 + Rg)), np.sqrt(2 / (Rg * (1 + Rg)))
check("1  dv1/dv2 = R (vt2+vc2)/(vc1+vt1)  (energy form)",
      np.allclose((vt1 - vc1) / (vc2 - vt2), Rg * (vt2 + vc2) / (vc1 + vt1), rtol=1e-10))
check("1  vc1+vt1 < R (vc2+vt2) for R>1", np.all(vc1 + vt1 < Rg * (vc2 + vt2)))
# 4: B(x) = (x-1)^2 (2x+1) has B' = 6x(x-1)
xb = sp.symbols("x")
check("4  B'(x) = 6x(x-1)", sp.expand(sp.diff((xb - 1)**2 * (2 * xb + 1), xb) - 6 * xb * (xb - 1)) == 0)
# 5: small-e estimates, dv_C ~ e v and dv_S ~ 2 e v sin(eta/2)
eb = sp.symbols("e", positive=True)
dvC = 2 * sp.sqrt(1 - eb) * (1 - sp.sqrt(1 - eb))
check("5  dv_C / (e v) -> 1 as e -> 0", sp.limit(dvC / eb, eb, 0) == 1)
# 8: LEO-GEO, the expansion parameter at the second burn is about 0.5 (not small)
r1g, r2g = 6678.137, 42164.0
vc2g = np.sqrt(mu / r2g)
vt2g = np.sqrt(2 * mu * r1g / (r2g * (r1g + r2g)))
epsg = 4 * vt2g * vc2g * np.sin(np.radians(28.5) / 2)**2 / (vc2g - vt2g)**2
print("     LEO-GEO epsilon at the second burn:", epsg)
check("8  LEO-GEO: epsilon at second burn ~ 0.5", 0.45 < epsg < 0.6)

print()
nf = sum(not ok for _, ok in RESULTS)
print(f"{len(RESULTS) - nf} PASS, {nf} FAIL")