"""Verification of every identity and number quoted in the CW note (cw-equations.qmd).

Part A is symbolic (sympy); Part B is numerical (two-body integration of both spacecraft).
Run:  python3 codes/cw-verify.py        (a few seconds)
"""
import numpy as np
import sympy as sp
from scipy.integrate import solve_ivp

ok = lambda label, cond: print(f"  [{'ok' if cond else 'FAIL'}] {label}")

# ======================================================================================= A. symbolic
print("A. symbolic")
t, n, mu, r = sp.symbols("t n mu r", positive=True)
x0, y0, xd0, yd0 = sp.symbols("x0 y0 xd0 yd0")
c, s = sp.cos(n*t), sp.sin(n*t)

# A1. rotating-frame kinematics, component route (Section 3.2)
x, y = sp.Function("x")(t), sp.Function("y")(t)
ihat, jhat = sp.Matrix([c, s]), sp.Matrix([-s, c])
acc = (x*ihat + y*jhat).diff(t, 2)
ok("inertial acceleration = (x''-2n y'-n^2 x) i + (y''+2n x'-n^2 y) j",
   sp.simplify(acc.dot(ihat) - (x.diff(t, 2) - 2*n*y.diff(t) - n**2*x)) == 0 and
   sp.simplify(acc.dot(jhat) - (y.diff(t, 2) + 2*n*x.diff(t) - n**2*y)) == 0)

# A2. gravity expansion (Section 4.1) and gradient tensor (Section 4.2)
eps, eta, e = sp.symbols("eps eta e")
D = 1 + 2*eps + eps**2 + eta**2
ser = sp.series((D**sp.Rational(-3, 2)).subs({eps: e*eps, eta: e*eta}), e, 0, 3).removeO().subs(e, 1)
ok("D^(-3/2) = 1 - 3 eps + 6 eps^2 - 3/2 eta^2 + ...", sp.expand(ser - (1 - 3*eps + 6*eps**2 - sp.Rational(3, 2)*eta**2)) == 0)
Rx = sp.expand(1 - (1 + eps)*ser)
Ry = sp.expand(-eta*ser)
ok("radial forcing / (n^2 r_t) = 2 eps - 3 eps^2 + 3/2 eta^2 (to 2nd order)",
   sp.expand(sp.series(Rx.subs({eps: e*eps, eta: e*eta}), e, 0, 3).removeO().subs(e, 1) - (2*eps - 3*eps**2 + sp.Rational(3, 2)*eta**2)) == 0)
ok("along-track forcing / (n^2 r_t) = -eta + 3 eps eta (to 2nd order)",
   sp.expand(sp.series(Ry.subs({eps: e*eps, eta: e*eta}), e, 0, 3).removeO().subs(e, 1) - (-eta + 3*eps*eta)) == 0)
Rv = sp.Matrix(sp.symbols("X1 X2 X3"))
avec = -mu*Rv/sp.sqrt(Rv.dot(Rv))**3
G = sp.simplify(avec.jacobian(Rv).subs({Rv[0]: r, Rv[1]: 0, Rv[2]: 0})*r**3/mu)
ok("gravity-gradient tensor = n^2 diag(2,-1,-1), trace 0", G == sp.diag(2, -1, -1))

# A3. the solution (Section 5)
xs = (4 - 3*c)*x0 + s/n*xd0 + 2*(1 - c)/n*yd0
ys = 6*(s - n*t)*x0 + y0 - 2*(1 - c)/n*xd0 + (4*s - 3*n*t)/n*yd0
ok("x(t), y(t) satisfy the CW equations", sp.simplify(xs.diff(t, 2) - 2*n*ys.diff(t) - 3*n**2*xs) == 0 and
   sp.simplify(ys.diff(t, 2) + 2*n*xs.diff(t)) == 0)
ok("initial conditions", [sp.simplify(q.subs(t, 0)) for q in (xs - x0, ys - y0, xs.diff(t) - xd0, ys.diff(t) - yd0)] == [0]*4)
ok("y' + 2 n x = y0' + 2 n x0", sp.simplify(ys.diff(t) + 2*n*xs - (yd0 + 2*n*x0)) == 0)
Jc = sp.simplify(((xs.diff(t))**2 + (ys.diff(t))**2)/2 - sp.Rational(3, 2)*n**2*xs**2)
ok("Jacobi integral constant", sp.simplify(Jc.diff(t)) == 0)

# A4. state transition matrix (Section 5.6-5.7)
X0 = sp.Matrix([x0, y0, xd0, yd0]); Xt = sp.Matrix([xs, ys, xs.diff(t), ys.diff(t)])
Phi = sp.simplify(Xt.jacobian(X0))
A = sp.Matrix([[0, 0, 1, 0], [0, 0, 0, 1], [3*n**2, 0, 0, 2*n], [0, 0, -2*n, 0]])
t1, t2 = sp.symbols("t1 t2", positive=True)
ok("Phi' = A Phi", sp.simplify(Phi.diff(t) - A*Phi) == sp.zeros(4, 4))
ok("Phi(t1+t2) = Phi(t2) Phi(t1)", sp.simplify(Phi.subs(t, t1 + t2) - Phi.subs(t, t2)*Phi.subs(t, t1)) == sp.zeros(4, 4))
ok("det Phi = 1", sp.simplify(Phi.det()) == 1)
lam = sp.symbols("lam")
ok("det(lam I - A) = lam^2 (lam^2 + n^2)", sp.expand((lam*sp.eye(4) - A).det() - lam**2*(lam**2 + n**2)) == 0)
v0, w = sp.Matrix([0, 1, 0, 0]), sp.Matrix([1, 0, 0, -3*n/2])
ok("A v = 0 and A w = -(3n/2) v (Jordan block)", A*v0 == sp.zeros(4, 1) and A*w == -sp.Rational(3, 2)*n*v0)
tau = sp.symbols("tau")
ok("det Phi_rv = [8(1-cos tau) - 3 tau sin tau]/n^2",
   sp.simplify(Phi[0:2, 2:4].det().subs(t, tau/n) - (8*(1 - sp.cos(tau)) - 3*tau*sp.sin(tau))/n**2) == 0)

# A5. guiding centre and epicycle form (Section 6.2)
xc, al, be = 4*x0 + 2*yd0/n, 3*x0 + 2*yd0/n, xd0/n
ok("x = x_c - alpha cos nt + beta sin nt", sp.simplify(xs - (xc - al*c + be*s)) == 0)
ok("y = y0 - 2 beta - (3/2) n x_c t + 2 alpha sin nt + 2 beta cos nt",
   sp.simplify(ys - (y0 - 2*be - sp.Rational(3, 2)*n*xc*t + 2*al*s + 2*be*c)) == 0)

# ======================================================================================= B. numerical
print("B. numerical")
mu_, RE = 398600.4418, 6378.137
a = RE + 420.0
n_ = np.sqrt(mu_/a**3); T = 2*np.pi/n_; v = a*n_
print(f"  a={a:.3f} km  n={n_:.6e} rad/s  T={T/60:.4f} min  v={v:.5f} km/s  g={mu_/a**2*1e3:.3f} m/s^2")
print(f"  tidal 3 n^2 x at 10 km = {3*n_**2*1e4:.4f} m/s^2   (ratio to g: {3*n_**2*1e4/(mu_/a**2*1e3):.2e})")
print(f"  mutual pull of a 4e5 kg target at 10 km: {6.674e-11*4e5/1e8:.1e} m/s^2")

def stm(tt):
    cc, ss, ta = np.cos(n_*tt), np.sin(n_*tt), n_*tt
    return np.array([[4 - 3*cc, 0, ss/n_, 2*(1 - cc)/n_], [6*(ss - ta), 1, -2*(1 - cc)/n_, (4*ss - 3*ta)/n_],
                     [3*n_*ss, 0, cc, 2*ss], [-6*n_*(1 - cc), 0, -2*ss, 4*cc - 3]])

def rhs(_, u):
    rr = np.hypot(u[0], u[1]); return [u[2], u[3], -mu_*u[0]/rr**3, -mu_*u[1]/rr**3]

def exact(X0, ts):
    """Integrate target and chaser inertially; return the offset in the LVLH frame."""
    x_, y_, xd_, yd_ = X0
    rt, vt = np.array([a, 0.0]), np.array([0.0, v])
    rc, vc = rt + [x_, y_], vt + [xd_, yd_] + n_*np.array([-y_, x_])
    kw = dict(t_eval=ts, rtol=1e-12, atol=1e-10)
    st = solve_ivp(rhs, [0, ts[-1]], [*rt, *vt], **kw); sc = solve_ivp(rhs, [0, ts[-1]], [*rc, *vc], **kw)
    out = []
    for k, tt in enumerate(ts):
        cc, ss = np.cos(n_*tt), np.sin(n_*tt); d = sc.y[:2, k] - st.y[:2, k]
        out.append([cc*d[0] + ss*d[1], -ss*d[0] + cc*d[1]])
    return np.array(out)

Dm = np.diag([1, 1, 1/n_, 1/n_])
print("  Phi(T/2) in the scaled variables (x, y, xdot/n, ydot/n):")
print(np.array2string(Dm @ stm(T/2) @ np.linalg.inv(Dm), precision=4, suppress_small=True))
dv = 1e-3
print(f"  Example 1 (1 m/s prograde): dv/n={dv/n_:.5f} km, 4dv/n={4*dv/n_:.4f}, 3 pi dv/n={3*np.pi*dv/n_:.4f}, 6 pi dv/n={6*np.pi*dv/n_:.4f}")
da = 2*dv/n_; dT = 1.5*T*da/a
print(f"    Kepler: da={da:.4f} km, dT={dT:.4f} s, lag v*dT={v*dT:.4f} km")
vn = v + dv; an = -mu_/(2*(vn**2/2 - mu_/a)); Tn = 2*np.pi*np.sqrt(an**3/mu_)
print(f"    exact orbit: da={an - a:.5f} km, dT={Tn - T:.5f} s, lag={v*(Tn - T):.4f} km")
ts = np.array([T/2, T]); ex = exact([0, 0, 0, dv], ts)
cw = np.array([stm(tt) @ [0, 0, 0, dv] for tt in ts])[:, :2]
print("    (x, y) at T/2, T  CW   :", np.round(cw, 4).tolist())
print("    (x, y) at T/2, T  exact:", np.round(ex, 4).tolist())
print(f"  Co-circular drift: 3/2 n * 1 km = {1.5*n_*1e3:.4f} m/s ; 3 pi = {3*np.pi:.4f} km per orbit")
print(f"  Radial burn: semi-axes dv/n={dv/n_:.4f} km, 2 dv/n={2*dv/n_:.4f} km")

# x_c = delta a: exact energy difference against 2 c / n
X0 = np.array([0.7e-3, 0.0, 0.3e-6, 0.5e-9])                       # small offsets (km, km/s)
rc_ = np.array([a + X0[0], X0[1]]); vc_ = np.array([0, v]) + [X0[2], X0[3]] + n_*np.array([-X0[1], X0[0]])
Ec = vc_ @ vc_/2 - mu_/np.hypot(*rc_); ac = -mu_/(2*Ec)
print(f"  x_c = 4 x0 + 2 yd0/n = {4*X0[0] + 2*X0[3]/n_:.6e} km ;  exact da = {ac - a:.6e} km")

print("  Validity (radial kick -> 2:1 ellipse of along-track extent d):")
ts = np.array([T/2, T, 2*T, 3*T])
for d in (1, 10, 100, 1000):
    X0 = np.array([0, 0, n_*d/4, 0]); ex = exact(X0, ts); cw = np.array([stm(tt) @ X0 for tt in ts])[:, :2]
    err = np.linalg.norm(ex - cw, axis=1)*1e3
    print(f"    d={d:5d} km  error after T/2, T, 2T, 3T [m] = {np.round(err, 3).tolist()}   err(T)/(d^2/a) = {err[1]*1e-3/(d**2/a):.3f}")
print(f"  Missing second-order delta a of the radial kick: 3 pi/16 = {3*np.pi/16:.4f}  (numerical coefficient above: 0.589)")
print(f"  A chaser on the same circle, 10 km and 100 km ahead, sits below the tangent by {10**2/(2*a)*1e3:.2f} m and {100**2/(2*a):.3f} km")