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

Random states are converted to elements with a plain textbook algorithm (rv2coe below,
with the three quadrant tests of the Orbital Elements deck), and each closed form in the
problem set is compared against it.
Run:  python codes/elements-analytic-verify.py      (prints PASS/FAIL per claim)
"""
import numpy as np
import sympy as sp
from scipy.integrate import quad

RESULTS = []
mu = 398600.4418
rng = np.random.default_rng(7)
K = np.array([0.0, 0.0, 1.0])


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


def wrap(x):
    return np.mod(x, 2 * np.pi)


def angdiff(x, y):
    return abs((x - y + np.pi) % (2 * np.pi) - np.pi)


def rv2coe(r, v):
    h = np.cross(r, v)
    n = np.cross(K, h)
    e = np.cross(v, h) / mu - r / np.linalg.norm(r)
    rn, hn, nn, en = map(np.linalg.norm, (r, h, n, e))
    eps = v @ v / 2 - mu / rn
    a = -mu / (2 * eps)
    i = np.arccos(h[2] / hn)
    Om = np.arccos(n[0] / nn)
    if n[1] < 0:
        Om = 2 * np.pi - Om
    w = np.arccos(np.clip(n @ e / (nn * en), -1, 1))
    if e[2] < 0:
        w = 2 * np.pi - w
    f = np.arccos(np.clip(e @ r / (en * rn), -1, 1))
    if r @ v < 0:
        f = 2 * np.pi - f
    return a, en, i, Om, w, f


def coe2rv(a, e, i, Om, w, f):
    p = a * (1 - e * e)
    r_pf = p / (1 + e * np.cos(f)) * np.array([np.cos(f), np.sin(f), 0])
    v_pf = np.sqrt(mu / p) * np.array([-np.sin(f), e + np.cos(f), 0])
    cO, sO, ci, si, cw, sw = np.cos(Om), np.sin(Om), np.cos(i), np.sin(i), np.cos(w), np.sin(w)
    Q = np.array([[cO * cw - sO * sw * ci, -cO * sw - sO * cw * ci, sO * si],
                  [sO * cw + cO * sw * ci, -sO * sw + cO * cw * ci, -cO * si],
                  [sw * si, cw * si, ci]])
    return Q @ r_pf, Q @ v_pf


def random_coe():
    return (rng.uniform(7000, 40000), rng.uniform(0.01, 0.9), rng.uniform(0.05, np.pi - 0.05),
            rng.uniform(0, 2 * np.pi), rng.uniform(0, 2 * np.pi), rng.uniform(0, 2 * np.pi))


samples = [random_coe() for _ in range(3000)]

# ---------------------------------------------------------------- 1. constants
ok1 = ok2 = True
for c in samples:
    r, v = coe2rv(*c)
    h = np.cross(r, v)
    e = np.cross(v, h) / mu - r / np.linalg.norm(r)
    eps = v @ v / 2 - mu / np.linalg.norm(r)
    ok1 &= abs(h @ e) / np.linalg.norm(h) < 1e-12
    ok2 &= abs(e @ e - (1 + 2 * eps * (h @ h) / mu**2)) < 1e-10
check("1  h . e = 0", ok1)
check("1  e^2 = 1 + 2 eps h^2 / mu^2", ok2)

# ---------------------------------------------------------------- 2. the local state (nu, gamma)
ok_a = ok_e = ok_f = ok_F = True
for c in samples:
    r, v = coe2rv(*c)
    a, e, *_, f = rv2coe(r, v)
    rn, vn = np.linalg.norm(r), np.linalg.norm(v)
    nu = rn * vn**2 / mu
    gam = np.arcsin((r @ v) / (rn * vn))
    ok_a &= abs(a - rn / (2 - nu)) < 1e-6
    ok_e &= abs(e**2 - (1 - nu * (2 - nu) * np.cos(gam)**2)) < 1e-10
    f2 = np.arctan2(nu * np.sin(gam) * np.cos(gam), nu * np.cos(gam)**2 - 1)
    ok_f &= angdiff(f, f2) < 1e-8
    # empty focus: at -2 a e_vec from the attracting focus; claim: P + (2a - r)(2 (rhat.vhat) vhat - rhat)
    h = np.cross(r, v)
    evec = np.cross(v, h) / mu - r / rn
    Fstar = -2 * a * evec
    rh, vh = r / rn, v / vn
    ok_F &= np.linalg.norm(Fstar - (r + (2 * a - rn) * (2 * (rh @ vh) * vh - rh))) < 1e-6 * a
check("2  a = r/(2 - nu)", ok_a)
check("2  e^2 = 1 - nu(2 - nu) cos^2 gamma", ok_e)
check("2  tan f = nu sin g cos g / (nu cos^2 g - 1)", ok_f)
check("2  empty focus = P + (2a - r)(2(rhat.vhat)vhat - rhat): PF reflected in the normal", ok_F)

# ---------------------------------------------------------------- 3. orbit averages
ok = True
for a, e in [(10000, 0.1), (20000, 0.6), (9000, 0.95)]:
    T = 2 * np.pi * np.sqrt(a**3 / mu)
    b = a * np.sqrt(1 - e * e)
    rE = lambda E: a * (1 - e * np.cos(E))
    # time average: dt = (rE/a) dE / n  ->  <g>_t = (1/2pi) int g(r(E)) (1 - e cos E) dE
    tav = lambda g: quad(lambda E: g(rE(E)) * (1 - e * np.cos(E)), 0, 2 * np.pi, limit=200)[0] / (2 * np.pi)
    fav = lambda g: quad(lambda f: g(a * (1 - e * e) / (1 + e * np.cos(f))), 0, 2 * np.pi, limit=200)[0] / (2 * np.pi)
    ok &= abs(tav(lambda r: 1 / r) * a - 1) < 1e-10
    ok &= abs(tav(lambda r: r) / (a * (1 + e * e / 2)) - 1) < 1e-10
    ok &= abs(tav(lambda r: r * r) / (a * a * (1 + 1.5 * e * e)) - 1) < 1e-10
    ok &= abs(tav(lambda r: 1 / r**2) * a * b - 1) < 1e-6
    ok &= abs(fav(lambda r: r) / b - 1) < 1e-10
    ok &= abs(fav(lambda r: 1 / r) * a * (1 - e * e) - 1) < 1e-10
    ok &= abs(tav(lambda r: mu * (2 / r - 1 / a)) * a / mu - 1) < 1e-10
    # fraction of time with r < a
    frac = quad(lambda E: (1 - e * np.cos(E)) * (rE(E) < a), 0, 2 * np.pi, points=[np.pi / 2, 3 * np.pi / 2], limit=200)[0] / (2 * np.pi)
    ok &= abs(frac - (0.5 - e / np.pi)) < 1e-9
check("3  <1/r>_t=1/a, <r>_t=a(1+e^2/2), <r^2>_t=a^2(1+3e^2/2), <1/r^2>_t=1/(ab), <r>_f=b, <1/r>_f=1/p, <v^2>_t=mu/a, P(r<a)=1/2-e/pi", ok)

# ---------------------------------------------------------------- 4. symmetries
okT = okM = True
for c in samples[:1000]:
    r, v = coe2rv(*c)
    a, e, i, Om, w, f = rv2coe(r, v)
    a2, e2, i2, Om2, w2, f2 = rv2coe(r, -v)
    okT &= (abs(a2 - a) < 1e-6 and abs(e2 - e) < 1e-10 and abs(i2 - (np.pi - i)) < 1e-9
            and angdiff(Om2, Om + np.pi) < 1e-9 and angdiff(w2, np.pi - w) < 1e-8 and angdiff(f2, -f) < 1e-8)
    M = np.diag([1, 1, -1.0])
    a3, e3, i3, Om3, w3, f3 = rv2coe(M @ r, M @ v)
    okM &= (abs(i3 - i) < 1e-9 and angdiff(Om3, Om + np.pi) < 1e-9 and angdiff(w3, w + np.pi) < 1e-8
            and angdiff(f3, f) < 1e-8)
check("4  v -> -v : (a, e, pi-i, Om+pi, pi-w, -f)", okT)
check("4  z -> -z : (a, e, i, Om+pi, w+pi, f)", okM)

# ---------------------------------------------------------------- 5. latitude, right ascension, azimuth
okd = oka = okb = True
for c in samples:
    r, v = coe2rv(*c)
    a, e, i, Om, w, f = rv2coe(r, v)
    u = w + f
    rn = np.linalg.norm(r)
    dec = np.arcsin(r[2] / rn)
    ra = np.arctan2(r[1], r[0])
    okd &= abs(np.sin(dec) - np.sin(i) * np.sin(u)) < 1e-12
    # tan(alpha - Omega) = cos i tan u, with the quadrant of (alpha - Omega) that of (cos u, cos i sin u)
    oka &= angdiff(ra - Om, np.arctan2(np.cos(i) * np.sin(u), np.cos(u))) < 1e-9
    # azimuth of the velocity, from north towards east
    east = np.cross(K, r)
    east /= np.linalg.norm(east)
    north = np.cross(r / rn, east)
    beta = np.arctan2(v @ east, v @ north)
    okb &= abs(np.cos(i) - np.cos(dec) * np.sin(beta)) < 1e-12
check("5  sin(dec) = sin i sin u", okd)
check("5  (alpha - Omega) = atan2(cos i sin u, cos u)", oka)
check("5  cos i = cos(dec) sin(azimuth)", okb)

# ---------------------------------------------------------------- 6. quadrant-free formulas
ok = True
for c in samples:
    r, v = coe2rv(*c)
    a, e, i, Om, w, f = rv2coe(r, v)
    h = np.cross(r, v)
    hn = np.linalg.norm(h)
    ev = np.cross(v, h) / mu - r / np.linalg.norm(r)
    Om2 = wrap(np.arctan2(h[0], -h[1]))
    i2 = np.arctan2(np.hypot(h[0], h[1]), h[2])
    w2 = wrap(np.arctan2(hn * ev[2], h[0] * ev[1] - h[1] * ev[0]))
    u2 = wrap(np.arctan2(hn * r[2], h[0] * r[1] - h[1] * r[0]))
    rn = np.linalg.norm(r)
    f2 = wrap(np.arctan2(hn * (r @ v), hn**2 - mu * rn))
    ok &= (angdiff(Om, Om2) < 1e-9 and abs(i - i2) < 1e-12 and angdiff(w, w2) < 1e-8
           and angdiff(w + f, u2) < 1e-8 and angdiff(f, f2) < 1e-8)
check("6  Om=atan2(hI,-hJ), i=atan2(sqrt(hI^2+hJ^2),hK), w=atan2(h eK, hI eJ - hJ eI), u=..., f=atan2(h r.v, h^2 - mu r)", ok)

# ---------------------------------------------------------------- 7. an out-of-plane impulse
ok = True
for c in samples[:500]:
    r, v = coe2rv(*c)
    a, e, i, Om, w, f = rv2coe(r, v)
    h = np.cross(r, v)
    hn = np.linalg.norm(h)
    W = h / hn
    dv = 1e-6
    a2, e2, i2, Om2, w2, f2 = rv2coe(r, v + dv * W)
    rn = np.linalg.norm(r)
    u = w + f
    di = rn * np.cos(u) / hn * dv
    dO = rn * np.sin(u) / (hn * np.sin(i)) * dv
    dw = -np.cos(i) * dO
    sc = rn / hn * dv / np.sin(i)
    ok &= (abs((i2 - i) - di) < 1e-4 * sc and abs(((Om2 - Om + np.pi) % (2 * np.pi) - np.pi) - dO) < 1e-4 * sc
           and abs(((w2 - w + np.pi) % (2 * np.pi) - np.pi) - dw) < 1e-4 * sc
           and abs(a2 - a) < 1e-6 * a and abs(e2 - e) < 1e-9 and angdiff(f2, f) < 1e-9)
check("7  di = r cos u dv/h, dOm = r sin u dv/(h sin i), dw = -cos i dOm, a, e, f unchanged (1st order)", ok)

# ---------------------------------------------------------------- 8. near-circular orbits
def kepler(M, e):
    E = M.copy()
    for _ in range(50):
        E = E - (E - e * np.sin(E) - M) / (1 - e * np.cos(E))
    return E


ok = True
M = np.linspace(0, 2 * np.pi, 721)
for e in [1e-3, 1e-4]:
    for w in [0.3, 2.0, 4.5]:
        a = 10000.0
        E = kepler(M, e)
        f = 2 * np.arctan2(np.sqrt(1 + e) * np.sin(E / 2), np.sqrt(1 - e) * np.cos(E / 2))
        r = a * (1 - e * np.cos(E))
        lam = w + M
        xi, eta = e * np.cos(w), e * np.sin(w)
        ok &= np.max(np.abs(r - a * (1 - xi * np.cos(lam) - eta * np.sin(lam)))) < 2 * a * e * e
        ok &= np.max(np.abs((w + f - lam - 2 * (xi * np.sin(lam) - eta * np.cos(lam)) + np.pi) % (2 * np.pi) - np.pi)) < 3 * e * e
check("8  r = a(1 - xi cos lam - eta sin lam) + O(e^2), u = lam + 2(xi sin lam - eta cos lam) + O(e^2)", ok)
# sensitivity of omega: d omega = (xi d eta - eta d xi)/e^2
xi, eta = 1e-3, 2e-4
d = 1e-9
w0 = np.arctan2(eta, xi)
w1 = np.arctan2(eta + d, xi - 0.5 * d)
check("8  d omega = (xi d eta - eta d xi)/e^2", abs((w1 - w0) - (xi * d - eta * (-0.5 * d)) / (xi**2 + eta**2)) < 1e-6 * abs(w1 - w0))

# ---------------------------------------------------------------- 9. time north of the equator
def frac_north(e, w):
    # u in (0, pi)  <=>  f in (-w, pi - w)
    # integrate dt/df = r^2/h  (normalised: dM/df = (1-e^2)^{3/2}/(1+e cos f)^2)
    g = lambda f: (1 - e * e)**1.5 / (1 + e * np.cos(f))**2
    return quad(g, -w, np.pi - w, limit=200, epsabs=1e-13)[0] / (2 * np.pi)


ok = True
for e in [0.1, 0.5, 0.74, 0.9]:
    near = (np.arccos(e) - e * np.sqrt(1 - e * e)) / np.pi
    ok &= abs(frac_north(e, 1.5 * np.pi) - (1 - near)) < 1e-8
    ok &= abs(frac_north(e, 0.5 * np.pi) - near) < 1e-8
    ok &= abs(frac_north(e, 0.0) - 0.5) < 1e-8
    ws = np.linspace(0, 2 * np.pi, 361)
    vals = np.array([frac_north(e, x) for x in ws])
    ok &= abs(ws[np.argmax(vals)] - 1.5 * np.pi) < 1e-9
check("9  fraction north: 1/2 (w=0), near-half (w=90), 1-near-half (w=270, the maximum)", ok)
print("     Molniya e=0.74: fraction north =", 1 - (np.arccos(0.74) - 0.74 * np.sqrt(1 - 0.74**2)) / np.pi)

# ---------------------------------------------------------------- 10. J2 secular rates
ci = 1 / np.sqrt(5)
check("10 critical inclination: tan i = 2", abs(np.tan(np.arccos(ci)) - 2) < 1e-12)
# Om_dot = -3/2 k cos i, w_dot = 3/4 k (5 cos^2 i - 1)
Omd = lambda i: -1.5 * np.cos(i)
wd = lambda i: 0.75 * (5 * np.cos(i)**2 - 1)
check("10 equatorial: Om_dot + w_dot = 3/2 k", abs(Omd(0) + wd(0) - 1.5) < 1e-15)
check("10 polar: Om_dot = 0, w_dot = -3/4 k", abs(Omd(np.pi / 2)) < 1e-15 and abs(wd(np.pi / 2) + 0.75) < 1e-15)
# the non-singular form: d(xi)/dt, d(eta)/dt at i=0 -- the pair (xi, eta) rotates at Om_dot + w_dot

# ---------------------------------------------------------------- steps added in the pedagogical rewrite
es = sp.symbols("e", positive=True)
check("9  d/de (arccos e - e sqrt(1-e^2)) = -2 sqrt(1-e^2)",
      sp.simplify(sp.diff(sp.acos(es) - es * sp.sqrt(1 - es**2), es) + 2 * sp.sqrt(1 - es**2)) == 0)
Eb = sp.symbols("E")
check("3  (1/2pi) int (1 - e cos E)^3 dE = 1 + 3e^2/2",
      sp.simplify(sp.integrate((1 - es * sp.cos(Eb))**3, (Eb, 0, 2 * sp.pi)) / (2 * sp.pi) - (1 + sp.Rational(3, 2) * es**2)) == 0)

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