"""Deterministic formula checks for Appendix G; run: python check_finite_range_bearing.py.

Synthetic sphere and known trajectories only: no constellation or MC validation.
"""
import numpy as np

RADIUS = 6_371_000.0
WAVELENGTH = 10.0
SPIN = np.array([[0., -7.292115e-5, 0.], [7.292115e-5, 0., 0.], [0., 0., 0.]])
TIMES = np.array([0., .13, .29, .51, .83])
DESIGN = np.column_stack((np.ones(len(TIMES)), TIMES))
ARRAY = np.array([[0., 0., 0.], [1., 0., 0.], [0., 1., 0.], [1., 1., 0.]])


def direction(x):
    a, b = x
    ca, sa, cb, sb = np.cos(a), np.sin(a), np.cos(b), np.sin(b)
    return np.array([cb*ca, cb*sa, sb]), np.array([
        [-cb*sa, -sb*ca], [cb*ca, -sb*sa], [0., cb]])


GROUND, _ = direction(np.array([.35, .42]))
GROUND *= RADIUS
SATELLITES = np.array([GROUND*1.10 + [90_000., -80_000., 40_000.],
                      GROUND*1.11 + [-210_000., 230_000., -60_000.]])
VELOCITIES = np.array([[1100., 6800., -1400.], [-1800., 6200., 2700.]])
ACCELERATIONS = np.array([[-4., -3., -1.], [-3., -2., -2.]])


def bearing(p, satellite):
    e = p - satellite
    e /= np.linalg.norm(e)
    return np.array([np.arctan2(e[1], e[0]), np.arcsin(e[2])])


def intersection(x):
    e, E = direction(x)
    s = SATELLITES[0]
    disc = (s@e)**2 - (s@s - RADIUS**2)
    assert disc > 0, "Ray is tangent or misses the surface"
    d = -s@e - np.sqrt(disc)
    assert d > 0
    p = s + d*e
    n = p / RADIUS
    assert n@e < 0, "Expected near visible branch"
    P = d*(np.eye(3) - np.outer(e, n)/(n@e))@E
    return p, P


def geometry(x, t, ell):
    p0, P0 = intersection(x)
    angle = 7.292115e-5*t
    c, s = np.cos(angle), np.sin(angle)
    B = np.array([[c, -s, 0.], [s, c, 0.], [0., 0., 1.]])
    p, P = B@p0, B@P0
    v = VELOCITIES[ell] + t*ACCELERATIONS[ell] - SPIN@p
    a = ACCELERATIONS[ell] - SPIN@SPIN@p
    Vx, Ax = -SPIN@P, -SPIN@SPIN@P
    rvec = SATELLITES[ell] + t*VELOCITIES[ell] + .5*t*t*ACCELERATIONS[ell] - p
    r = np.linalg.norm(rvec)
    w = rvec/r
    S = np.eye(3) - np.outer(w, w)
    W = -S@P/r
    wd = S@v/r
    rx = -w@P
    WD = np.column_stack([
        -(np.outer(W[:, j], w) + np.outer(w, W[:, j]))@v/r
        + S@Vx[:, j]/r - wd*rx[j]/r for j in range(2)
    ])
    b = -v@w/WAVELENGTH
    bd = -(a@w + v@wd)/WAVELENGTH
    bx = -(w@Vx + v@W)/WAVELENGTH
    bdx = -(w@Ax + a@W + wd@Vx + v@WD)/WAVELENGTH
    return w, W, wd, WD, b, bx, bd, bdx


def coefficients(x, ell):
    data = [geometry(x, t, ell) for t in TIMES]
    ref = [geometry(x, t, 0) for t in TIMES]
    c = np.linalg.lstsq(DESIGN, [a[4]-b[4] for a, b in zip(data, ref)], rcond=None)[0]
    Cx = np.linalg.lstsq(DESIGN, [a[5]-b[5] for a, b in zip(data, ref)], rcond=None)[0]
    return c, Cx


def mean_and_jacobian(x, ref_coeff, ell):
    c, Cx = coefficients(x, ell)
    f, rate = ref_coeff + c
    w, W = geometry(x, 0., ell)[:2]
    steering = np.exp(-2j*np.pi*ARRAY@w/WAVELENGTH)
    Ax = (-2j*np.pi/WAVELENGTH)*steering[:, None]*(ARRAY@W)
    q = np.exp(2j*np.pi*(f*TIMES + .5*rate*TIMES**2))
    qx = q[:, None]*2j*np.pi*(TIMES[:, None]*Cx[0] + .5*TIMES[:, None]**2*Cx[1])
    Jx = np.column_stack([np.kron(q, Ax[:, j])+np.kron(qx[:, j], steering) for j in range(2)])
    Jf = np.kron(2j*np.pi*TIMES*q, steering)
    Jr = np.kron(1j*np.pi*TIMES**2*q, steering)
    return np.kron(q, steering), np.column_stack((Jx, Jf, Jr))


def finite_difference(f, x, step=2e-7):
    return np.stack([(f(x+step*e)-f(x-step*e))/(2*step) for e in np.eye(len(x))], axis=-1)


def check(name, actual, expected, tolerance=2e-7):
    error = np.linalg.norm(actual-expected)/max(1., np.linalg.norm(expected))
    assert error < tolerance, (name, error, tolerance)
    print(f"{name}: relative error {error:.3g}")


def main():
    x = bearing(GROUND, SATELLITES[0])
    p, P = intersection(x)
    check("ray intersection", p, GROUND)
    check("intersection Jacobian", P, finite_difference(lambda x: intersection(x)[0], x))
    latlon = np.array([.35, .42])
    _, Pll = direction(latlon)
    Pll *= RADIUS
    e = (p-SATELLITES[0])/np.linalg.norm(p-SATELLITES[0])
    K = np.array([[-e[1]/(e[0]**2+e[1]**2), e[0]/(e[0]**2+e[1]**2), 0.],
                  [0., 0., 1/np.sqrt(1-e[2]**2)]])
    H = K@(np.eye(3)-np.outer(e, e))@Pll/np.linalg.norm(p-SATELLITES[0])
    check("bearing/latlon tangent identity", P@H, Pll)
    for ell in range(2):
        w, W, wd, WD, b, bx, bd, bdx = geometry(x, .41, ell)
        for name, index, jac in [("LOS", 0, W), ("LOS time derivative", 2, WD),
                                 ("Doppler", 4, bx), ("rate", 6, bdx)]:
            check(name, jac, finite_difference(lambda x: geometry(x, .41, ell)[index], x))
        step = 1e-3
        check("rate time derivative", bd,
              (geometry(x, .41+step, ell)[4]-geometry(x, .41-step, ell)[4])/(2*step))
    c, Cx = coefficients(x, 1)
    check("OLS coefficient Jacobian", Cx, finite_difference(lambda x: coefficients(x, 1)[0], x))
    z = np.r_[x, 20., .8]
    means, jac = zip(*(mean_and_jacobian(x, z[2:], ell) for ell in range(2)))
    J = np.vstack(jac)
    check("CP full mean Jacobian", J,
          finite_difference(lambda z: np.concatenate([mean_and_jacobian(z[:2], z[2:], ell)[0]
                                                      for ell in range(2)]), z), 2e-6)
    T = np.zeros((4, 4))
    T[:2, :2] = H
    T[2:, 2:] = np.eye(2)
    Jll = finite_difference(lambda z: np.concatenate([
        mean_and_jacobian(bearing(RADIUS*direction(z[:2])[0], SATELLITES[0]), z[2:], ell)[0]
        for ell in range(2)]), np.r_[latlon, z[2:]], step=1e-8)
    check("direct versus chain Jacobian", J@T, Jll, 2e-6)
    A = np.vstack((J.real, J.imag))*np.sqrt(2)
    nuisance = np.zeros((len(A), 4))
    for ell in range(2):
        gain = np.concatenate([means[i] if i == ell else np.zeros_like(means[i]) for i in range(2)])
        nuisance[:, 2*ell] = np.r_[gain.real, gain.imag]*np.sqrt(2)
        nuisance[:, 2*ell+1] = np.r_[-gain.imag, gain.real]*np.sqrt(2)
    projected = (np.eye(len(A))-nuisance@np.linalg.pinv(nuisance))@A
    info = projected.T@projected
    assert np.linalg.matrix_rank(projected) == 4
    cov = np.linalg.inv(info)
    check("full covariance reparameterization", T@np.linalg.inv(T.T@info@T)@T.T, cov)
    delta = .19
    shifted = TIMES-delta
    for ell in range(2):
        f, rate = z[2:]+coefficients(x, ell)[0]
        phase0 = 2*np.pi*(f*TIMES+.5*rate*TIMES**2)
        phase1 = 2*np.pi*(f*delta+.5*rate*delta**2
                          +(f+delta*rate)*shifted+.5*rate*shifted**2)
        check("epoch phase identity", np.exp(1j*phase0), np.exp(1j*phase1))
    def switch(z):
        p, _ = intersection(z[:2])
        return np.r_[bearing(p, SATELLITES[1]), z[2:]+coefficients(z[:2], 1)[0]]
    Tref = np.eye(4)
    Tref[:2, :2] = finite_difference(lambda x: bearing(intersection(x)[0], SATELLITES[1]), x)
    Tref[2:, :2] = Cx
    check("reference-change full Jacobian", Tref, finite_difference(switch, z), 2e-6)
    print("PASS: synthetic local formula checks; code-side geometry and CP/IP acceptance still required.")


if __name__ == "__main__":
    main()