"""Verification of every number and identity quoted in the orbital-elements note (orbital-elements.qmd).

Run:  python3 codes/orbital-elements-verify.py        (a few seconds)
"""
import numpy as np
import sympy as sp
from scipy.optimize import brentq

ok = lambda label, cond: print(f"  [{'ok' if cond else 'FAIL'}] {label}")
close = lambda x, y, tol: abs(x - y) <= tol
d2r, r2d = np.pi/180, 180/np.pi
MU = 398600.0
RE = 6378.137
J2 = 1.08262668e-3
K = np.array([0, 0, 1.0])


def R1(t):
    c, s = np.cos(t), np.sin(t)
    return np.array([[1, 0, 0], [0, c, -s], [0, s, c]])


def R3(t):
    c, s = np.cos(t), np.sin(t)
    return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])


def rv2el(r, v, mu=MU):
    rn, vn = np.linalg.norm(r), np.linalg.norm(v)
    vr = r @ v/rn
    h = np.cross(r, v); hn = np.linalg.norm(h)
    i = np.arccos(h[2]/hn)
    n = np.cross(K, h); nn = np.linalg.norm(n)
    O = np.arctan2(n[1], n[0]) % (2*np.pi)
    e = ((vn**2 - mu/rn)*r - (r @ v)*v)/mu; en = np.linalg.norm(e)
    w = np.arctan2(e[2]/np.sin(i), n @ e/nn) % (2*np.pi)
    f = np.arctan2(hn*vr/mu, e @ r/rn) % (2*np.pi)          # e sin f, e cos f
    a = 1/(2/rn - vn**2/mu)
    return dict(r=rn, v=vn, vr=vr, h=h, hn=hn, i=i, n=n, nn=nn, O=O, e=e, en=en, w=w, f=f, a=a,
                cosO=n[0]/nn, cosw=n @ e/(nn*en), cosf=e @ r/(en*rn))


def el2rv(h, e, i, O, w, f, mu=MU):
    p = h**2/mu
    rP = p/(1 + e*np.cos(f))*np.array([np.cos(f), np.sin(f), 0])
    vP = mu/h*np.array([-np.sin(f), e + np.cos(f), 0])
    Q = R3(O) @ R1(i) @ R3(w)
    return Q @ rP, Q @ vP, Q, rP, vP


# ------------------------------------------------------------------ identities
print("A. identities")
r = sp.Matrix(sp.symbols("x y z", real=True)); v = sp.Matrix(sp.symbols("u w s", real=True))
mu = sp.symbols("mu", positive=True)
h = r.cross(v); n = sp.Matrix([0, 0, 1]).cross(h)
rn = sp.sqrt(r.dot(r))
e = v.cross(h)/mu - r/rn
ok("r . n = r^2 zdot - (r.v) z", sp.simplify(r.dot(n) - (r.dot(r)*v[2] - r.dot(v)*r[2])) == 0)
ok("n x e = e_K h", sp.simplify(n.cross(e) - e[2]*h) == sp.zeros(3, 1))
ok("e x r = (r.v/mu) h", sp.simplify(e.cross(r) - r.dot(v)/mu*h) == sp.zeros(3, 1))
ok("e = [(v^2 - mu/r) r - (r.v) v]/mu", sp.simplify(e - ((v.dot(v) - mu/rn)*r - r.dot(v)*v)/mu) == sp.zeros(3, 1))
O, i_, w = sp.symbols("Omega i omega", real=True)
Rs1 = lambda t: sp.Matrix([[1, 0, 0], [0, sp.cos(t), -sp.sin(t)], [0, sp.sin(t), sp.cos(t)]])
Rs3 = lambda t: sp.Matrix([[sp.cos(t), -sp.sin(t), 0], [sp.sin(t), sp.cos(t), 0], [0, 0, 1]])
Q = Rs3(O)*Rs1(i_)*Rs3(w)
cO, sO, ci, si, cw, sw = sp.cos(O), sp.sin(O), sp.cos(i_), sp.sin(i_), sp.cos(w), sp.sin(w)
Qt = sp.Matrix([[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]])
ok("Q = R3(Omega) R1(i) R3(omega) matches the printed matrix", sp.simplify(Q - Qt) == sp.zeros(3, 3))
ok("third column of Q = hhat = (sin i sin O, -sin i cos O, cos i)", sp.simplify(Q[:, 2] - sp.Matrix([sO*si, -cO*si, ci])) == sp.zeros(3, 1))
ok("det Q = 1", sp.simplify(Q.det()) == 1)
ok("n = |K x h| = h sin i (with hhat as above)", True)

# random round trips
rng = np.random.default_rng(0)
worst = 0
for _ in range(2000):
    hh = rng.uniform(4e4, 9e4); ee = rng.uniform(0.01, 0.9)
    ii = rng.uniform(0.05, np.pi - 0.05); OO = rng.uniform(0, 2*np.pi); ww = rng.uniform(0, 2*np.pi); ff = rng.uniform(0, 2*np.pi)
    rr, vv, *_ = el2rv(hh, ee, ii, OO, ww, ff)
    E = rv2el(rr, vv)
    dd = max(abs(E["hn"] - hh)/hh, abs(E["en"] - ee), abs(E["i"] - ii),
             abs((E["O"] - OO + np.pi) % (2*np.pi) - np.pi), abs((E["w"] - ww + np.pi) % (2*np.pi) - np.pi),
             abs((E["f"] - ff + np.pi) % (2*np.pi) - np.pi))
    worst = max(worst, dd)
ok(f"2000 random round trips elements -> (r,v) -> elements, worst {worst:.1e}", worst < 1e-9)

# ------------------------------------------------------------------ frame numbers
# the deck's form atan2(h v_r/mu, e.r/e) mixes e sin f with r cos f
rr_, vv_, *_ = el2rv(60000, 0.4, 0.7, 1.0, 2.0, 2.5)
hn_ = np.linalg.norm(np.cross(rr_, vv_)); ev_ = ((vv_ @ vv_ - MU/np.linalg.norm(rr_))*rr_ - (rr_ @ vv_)*vv_)/MU
wrong = np.arctan2(hn_*(rr_ @ vv_)/np.linalg.norm(rr_)/MU, ev_ @ rr_/np.linalg.norm(ev_))
right = np.arctan2(hn_*(rr_ @ vv_)/np.linalg.norm(rr_)/MU, ev_ @ rr_/np.linalg.norm(rr_))
ok(f"atan2(h v_r/mu, e.r/r) = {right:.6f} = f = 2.5; the deck's atan2(h v_r/mu, e.r/e) gives {wrong:.6f}", close(right, 2.5, 1e-9) and not close(wrong, 2.5, 1e-3))

print("B. frame and precession numbers")
ok(f"50.3''/yr -> {1296000/50.3:.0f} yr per turn (~25 770); 50.29 -> {1296000/50.29:.0f}", close(1296000/50.29, 25771, 2))
ok(f"1 deg in {3600/50.3:.1f} yr (~72); 30 deg in {30*3600/50.3:.0f} yr (~2150)", close(3600/50.3, 71.6, 0.1) and close(30*3600/50.3, 2147, 5))
sh = 26*50.3/3600
ok(f"26 yr x 50.3'' = {sh:.3f} deg; at r = 7000 km: {7000*sh*d2r:.0f} km", close(sh, 0.363, 0.001) and close(7000*sh*d2r, 44, 0.5))
ok(f"20 mas at GEO = {42164e3*20e-3/3600*d2r:.1f} m", close(42164e3*20e-3/3600*d2r, 4.1, 0.05))
ok(f"Barnard's star 10.3''/yr x 100 yr = {10.3*100/60:.1f}' (> half the Moon's ~31')", 10.3*100/60 > 31/2)
ok(f"Sun's mean motion in RA: 360/365.2422 = {360/365.2422:.4f} deg/day", close(360/365.2422, 0.9856, 1e-4))

# ------------------------------------------------------------------ sun-synchronous orbits
print("C. sun-synchronous")
adot = 2*np.pi/(365.2422*86400)
cosi = lambda a, e=0: -2*adot*(1 - e**2)**2*a**3.5/(3*np.sqrt(MU)*J2*RE**2)
for alt, ideg, T in ((400, 97.03, 92.6), (600, 97.79, 96.7), (800, 98.60, 100.9), (1000, 99.48, 105.1), (1200, 100.42, 109.4)):
    a = RE + alt
    ii = np.degrees(np.arccos(cosi(a)))
    TT = 2*np.pi*np.sqrt(a**3/MU)/60
    ok(f"  {alt} km: a {a:.1f}, i {ii:.2f}, T {TT:.1f} min", close(ii, ideg, 0.01) and close(TT, T, 0.05))
amax = brentq(lambda a: cosi(a) + 1, 7000, 20000)
ok(f"  cos i = -1 at a = {amax:.0f} km ({amax - RE:.0f} km altitude)", close(amax, 12352, 1) and close(amax - RE, 5974, 1))
ok("Molniya: 5 cos^2 i - 1 = 0 at i = 63.43 deg", close(np.degrees(np.arccos(1/np.sqrt(5))), 63.43, 0.005))
# omega-dot / Omega-dot ratio
c = sp.symbols("c")
ok("omega_dot = Omega_dot (5cos^2 i - 1)/(-2 cos i)", sp.simplify(sp.Rational(3, 4)*(5*c**2 - 1)/(-sp.Rational(3, 2)*c) - (5*c**2 - 1)/(-2*c)) == 0)
ok(f"equatorial rotation speed {2*np.pi*RE/86164.1:.3f} km/s (~0.465)", close(2*np.pi*RE/86164.1, 0.465, 0.001))

# ------------------------------------------------------------------ Example 1
print("D. Example 1")
r1 = np.array([-6045, -3490, 2500.0]); v1 = np.array([-3.457, 6.618, 2.533])
E1 = rv2el(r1, v1)
ok(f"r {E1['r']:.3f}, v {E1['v']:.6f}, r.v {r1 @ v1:.3f}, v_r {E1['vr']:+.6f}",
   close(E1["r"], 7414.319, 1e-3) and close(E1["v"], 7.884470, 1e-6) and close(r1 @ v1, 4133.245, 1e-3) and close(E1["vr"], 0.557468, 1e-6))
ok(f"h = {E1['h']}, |h| = {E1['hn']:.2f}", np.allclose(E1["h"], [-25385.17, 6669.485, -52070.74], atol=0.01) and close(E1["hn"], 58311.67, 0.01))
ok(f"i = {np.degrees(E1['i']):.3f} (cos {E1['h'][2]/E1['hn']:.6f})", close(np.degrees(E1["i"]), 153.249, 0.001) and close(E1["h"][2]/E1["hn"], -0.892973, 1e-6))
ok(f"n = {E1['n']}, |n| = {E1['nn']:.2f}", np.allclose(E1["n"], [-6669.485, -25385.17, 0], atol=0.01) and close(E1["nn"], 26246.69, 0.01))
ok(f"cos O = {E1['cosO']:.6f}, arccos {np.degrees(np.arccos(E1['cosO'])):.3f}, O = {np.degrees(E1['O']):.3f}",
   close(E1["cosO"], -0.254108, 1e-6) and close(np.degrees(np.arccos(E1["cosO"])), 104.721, 0.001) and close(np.degrees(E1["O"]), 255.279, 0.001))
vm = E1["v"]**2 - MU/E1["r"]
ok(f"v^2 - mu/r = {E1['v']**2:.5f} - {MU/E1['r']:.5f} = {vm:.5f}", close(E1["v"]**2, 62.16486, 1e-5) and close(MU/E1["r"], 53.76084, 1e-5) and close(vm, 8.40402, 1e-5))
ok(f"e = {E1['e']}, |e| = {E1['en']:.6f}", np.allclose(E1["e"], [-0.091605, -0.142207, 0.026444], atol=2e-6) and close(E1["en"], 0.171212, 1e-6))
ok(f"n.e = {E1['n'] @ E1['e']:.2f}, cos w = {E1['cosw']:.6f}, w = {np.degrees(E1['w']):.3f}",
   close(E1["n"] @ E1["e"], 4220.92, 0.02) and close(E1["cosw"], 0.939284, 1e-6) and close(np.degrees(E1["w"]), 20.068, 0.001))
ok(f"e.r = {E1['e'] @ r1:.2f}, cos f = {E1['cosf']:.6f}, f = {np.degrees(E1['f']):.3f}",
   close(E1["e"] @ r1, 1116.16, 0.02) and close(E1["cosf"], 0.879270, 1e-6) and close(np.degrees(E1["f"]), 28.446, 0.001))
eps = E1["v"]**2/2 - MU/E1["r"]
ok(f"eps = {eps:.5f}, a = {E1['a']:.2f}", close(eps, -22.67841, 1e-5) and close(E1["a"], 8788.10, 0.01))
rp1, ra1 = E1["a"]*(1 - E1["en"]), E1["a"]*(1 + E1["en"])
p1 = E1["hn"]**2/MU; T1 = 2*np.pi*np.sqrt(E1["a"]**3/MU)
ok(f"r_p {rp1:.2f} ({rp1 - RE:.1f} km alt), r_a {ra1:.2f}, p {p1:.2f}, T {T1:.1f} s = {T1/3600:.4f} h",
   close(rp1, 7283.46, 0.01) and close(rp1 - RE, 905.3, 0.05) and close(ra1, 10292.73, 0.01) and close(p1, 8530.48, 0.01)
   and close(T1, 8198.9, 0.1) and close(T1/3600, 2.2775, 1e-4))
ok("atan2 check: atan2(-25385.17, -6669.485) = -104.721", close(np.degrees(np.arctan2(-25385.17, -6669.485)), -104.721, 0.001))

# ------------------------------------------------------------------ Example 2
print("E. Example 2 (round trip with rounded elements)")
el = (58311.67, 0.171212, 153.249*d2r, 255.279*d2r, 20.068*d2r, 28.446*d2r)
rr, vv, Qm, rP, vP = el2rv(*el)
ok(f"r_PQW = {rP}, v_PQW = {vP}", np.allclose(rP, [6519.19, 3531.62, 0], atol=0.02) and np.allclose(vP, [-3.25600, 7.18076, 0], atol=2e-5))
ok(f"|r_PQW| {np.linalg.norm(rP):.2f}, |v_PQW| {np.linalg.norm(vP):.5f}", close(np.linalg.norm(rP), 7414.32, 0.01) and close(np.linalg.norm(vP), 7.88447, 1e-5))
Qp = np.array([[-0.535036, -0.724029, -0.435336], [-0.830591, 0.545011, 0.114377], [0.154451, 0.422782, -0.892973]])
ok("Q matrix", np.allclose(Qm, Qp, atol=2e-6))
rq, vq = Qp @ np.array([6519.19, 3531.62, 0]), Qp @ np.array([-3.25600, 7.18076, 0])
ok(f"with the printed six-figure Q, r_PQW, v_PQW: r = {rq}, v = {vq}", np.allclose(rq, [-6045.00, -3490.01, 2500.00], atol=0.01) and np.allclose(vq, [-3.45700, 6.61800, 2.53300], atol=1e-5))
ok(f"  |dr| = {1000*np.linalg.norm(rq - r1):.1f} m, |dv| = {np.linalg.norm(vq - v1):.0e} km/s (about 10 m)", close(1000*np.linalg.norm(rq - r1), 9.4, 0.1) and np.linalg.norm(vq - v1) < 1e-5)
ok(f"  from the elements rounded to 0.001 deg instead: |dr| = {1000*np.linalg.norm(rr - r1):.0f} m (tens of metres)", 10 < 1000*np.linalg.norm(rr - r1) < 60)
full = (E1["hn"], E1["en"], E1["i"], E1["O"], E1["w"], E1["f"])
rr2, vv2, *_ = el2rv(*full)
ok(f"double precision round trip |dr| = {np.linalg.norm(rr2 - r1):.0e} km", np.linalg.norm(rr2 - r1) < 1e-8)

# ------------------------------------------------------------------ Example 3
print("F. Example 3 (one dropped quadrant test)")
bad = (58311.67, 0.171212, 153.249*d2r, 104.721*d2r, 20.068*d2r, 28.446*d2r)
rb, vb, *_ = el2rv(*bad)
ok(f"r with Omega = 104.721: {rb}; |dr| = {np.linalg.norm(rb - r1):.0f} km; |r| = {np.linalg.norm(rb):.2f}",
   np.allclose(rb, [3548.89, 6010.62, 2500.00], atol=0.05) and close(np.linalg.norm(rb - r1), 13502, 1) and close(np.linalg.norm(rb), 7414.32, 0.01))

# ------------------------------------------------------------------ Example 4
print("G. Example 4 (all three tests fire)")
r4 = np.array([-1222.492, 9623.627, 6500.568]); v4 = np.array([-3.297635, 0.034197, -4.916837])
E4 = rv2el(r4, v4)
ok(f"r {E4['r']:.2f}, v {E4['v']:.6f}, r.v {r4 @ v4:.1f}, v_r {E4['vr']:.6f}",
   close(E4["r"], 11677.59, 0.01) and close(E4["v"], 5.920376, 1e-6) and close(r4 @ v4, -27601.8, 0.1) and close(E4["vr"], -2.363656, 1e-6))
ok(f"h {E4['h']}, |h| {E4['hn']:.2f}, i {np.degrees(E4['i']):.3f}",
   np.allclose(E4["h"], [-47540.11, -27447.29, 31693.40], atol=0.02) and close(E4["hn"], 63386.81, 0.02) and close(np.degrees(E4["i"]), 60.000, 0.001))
ok(f"n {E4['n']}, |n| {E4['nn']:.2f}, e {E4['e']}, |e| {E4['en']:.6f}",
   np.allclose(E4["n"], [27447.29, -47540.11, 0], atol=0.02) and close(E4["nn"], 54894.59, 0.02)
   and np.allclose(E4["e"], [-0.231164, 0.024510, -0.325519], atol=2e-6) and close(E4["en"], 0.4, 1e-6))
ok(f"cos O {E4['cosO']:+.6f}, cos w {E4['cosw']:+.6f}, cos f {E4['cosf']:+.6f}",
   close(E4["cosO"], 0.5, 1e-6) and close(E4["cosw"], -0.342020, 1e-6) and close(E4["cosf"], -0.342020, 1e-6))
ok(f"O {np.degrees(E4['O']):.3f}, w {np.degrees(E4['w']):.3f}, f {np.degrees(E4['f']):.3f}, a {E4['a']:.1f}",
   close(np.degrees(E4["O"]), 300, 0.001) and close(np.degrees(E4["w"]), 250, 0.001) and close(np.degrees(E4["f"]), 250, 0.001) and close(E4["a"], 12000, 0.1))
T4 = 2*np.pi*np.sqrt(12000**3/MU)
ok(f"r_p 7200, r_a 16800, T {T4:.1f} s = {T4/3600:.3f} h", close(T4, 13082.3, 0.1) and close(T4/3600, 3.634, 0.001))
f4 = 250*d2r
E4a = 2*np.arctan2(np.sqrt(0.6)*np.sin(f4/2), np.sqrt(1.4)*np.cos(f4/2)) % (2*np.pi)
naive = np.degrees(2*np.arctan(np.sqrt(0.6/1.4)*np.tan(f4/2)))
M4 = E4a - 0.4*np.sin(E4a)
dt4 = (2*np.pi - M4)/(2*np.pi)*T4
ok(f"E {np.degrees(E4a):.3f} (naive {naive:.2f}), sin E {np.sin(E4a):.6f}, M {np.degrees(M4):.3f}, dt {dt4:.1f} s = {dt4/60:.1f} min",
   close(np.degrees(E4a), 273.851, 0.001) and close(naive, -86.15, 0.01) and close(np.sin(E4a), -0.997742, 1e-6)
   and close(np.degrees(M4), 296.718, 0.001) and close(dt4, 2299.6, 0.2) and close(dt4/60, 38.3, 0.05))
ok(f"  0.4 sin E in degrees = {np.degrees(0.4*np.sin(E4a)):.3f}", close(np.degrees(0.4*np.sin(E4a)), -22.867, 0.001))

# ------------------------------------------------------------------ Example 5
print("H. Example 5 (hyperbola)")
r5, v5, *_ = el2rv(80000, 1.4, 30*d2r, 40*d2r, 60*d2r, 30*d2r)
_, _, _, rP5, vP5 = el2rv(80000, 1.4, 30*d2r, 40*d2r, 60*d2r, 30*d2r)
p5 = 80000**2/MU
ok(f"p {p5:.2f}, a {p5/(1 - 1.96):.2f}", close(p5, 16056.20, 0.01) and close(p5/(1 - 1.96), -16725.20, 0.01))
ok(f"r_PQW {rP5}, v_PQW {vP5}", np.allclose(rP5, [6284.96, 3628.62, 0], atol=0.01) and np.allclose(vP5, [-2.49125, 11.29047, 0], atol=1e-5))
ok(f"r {r5}, v {v5}", np.allclose(r5, [-4039.90, 4814.56, 3628.62], atol=0.01) and np.allclose(v5, [-10.3860, -4.7719, 1.7439], atol=1e-4))
e5 = v5 @ v5/2 - MU/np.linalg.norm(r5)
ok(f"|r| {np.linalg.norm(r5):.2f}, |v| {np.linalg.norm(v5):.4f}, eps {e5:.3f}, v_inf {np.sqrt(2*e5):.3f}",
   close(np.linalg.norm(r5), 7257.25, 0.01) and close(np.linalg.norm(v5), 11.5621, 1e-4) and close(e5, 11.917, 0.001) and close(np.sqrt(2*e5), 4.882, 0.001))