"""Deterministic formula checks for Appendices F/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 mixed_motion(x, t, ell):
    """Synthetic non-spin rotation: deliberately not an astronomical model."""
    p0, P0 = intersection(x)
    a, k = 7.292115e-5*t, 3e-5
    z = np.array([[np.cos(a), -np.sin(a), 0.],
                  [np.sin(a), np.cos(a), 0.], [0., 0., 1.]])
    a = k*t
    q = np.array([[1., 0., 0.], [0., np.cos(a), -np.sin(a)],
                  [0., np.sin(a), np.cos(a)]])
    generator = np.array([[0., 0., 0.], [0., 0., -k], [0., k, 0.]])
    B = z@q
    Bd = SPIN@B + B@generator
    C, Cd = SPIN@B, SPIN@Bd
    p, pd, vu = B@p0, Bd@p0, C@p0
    s = SATELLITES[ell] + t*VELOCITIES[ell] + .5*t*t*ACCELERATIONS[ell]
    vs = VELOCITIES[ell] + t*ACCELERATIONS[ell]
    r = np.linalg.norm(s-p)
    w = (s-p)/r
    S = np.eye(3)-np.outer(w, w)
    wd = S@(vs-pd)/r
    proxy = S@(vs-vu)/r
    W = -S@B@P0/r
    b = -(vs-vu)@w/WAVELENGTH
    bg = -(vs-pd)@w/WAVELENGTH
    bd = -((ACCELERATIONS[ell]-Cd@p0)@w + (vs-vu)@wd)/WAVELENGTH
    bx = -(-w@C@P0+(vs-vu)@W)/WAVELENGTH
    return p, pd, vu, w, wd, proxy, b, bg, bd, bx


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)
    # Appendix F: structured CP has real frame amplitudes and one phase per star.
    for selected in [[0], [0, 1]]:
        Jcp = np.vstack([jac[ell] for ell in selected])
        N = np.zeros((len(Jcp), len(selected)*(len(TIMES)+1)), dtype=complex)
        block_size = len(TIMES)*len(ARRAY)
        for star, ell in enumerate(selected):
            rows = slice(star*block_size, (star+1)*block_size)
            N[rows, star*(len(TIMES)+1)] = 1j*means[ell]
            for frame in range(len(TIMES)):
                block = slice(star*block_size+frame*len(ARRAY),
                              star*block_size+(frame+1)*len(ARRAY))
                N[block, star*(len(TIMES)+1)+frame+1] = means[ell][frame*len(ARRAY):(frame+1)*len(ARRAY)]
        real_N = np.vstack((N.real, N.imag))
        scores = np.vstack((Jcp.real, Jcp.imag))*np.sqrt(2)
        residual_scores = scores-real_N@np.linalg.lstsq(real_N, scores, rcond=None)[0]
        I = residual_scores.T@residual_scores
        K, coupling, rate_info = I[:3, :3], I[:3, 3], I[3, 3]
        CK = np.linalg.inv(K)
        CU = np.linalg.inv(K-np.outer(coupling, coupling)/rate_info)
        v = CK@coupling
        h = rate_info-coupling@v
        assert h > 0, "K/U covariance identities require identifiable rate"
        check("CP K/U full covariance update", CU-CK, np.outer(v, v)/h, 1e-9)
        q = coupling[:2]-K[:2, 2]*coupling[2]/K[2, 2]
        G = K[:2, :2]-np.outer(K[:2, 2], K[:2, 2])/K[2, 2]
        check("partial angle-rate coupling", v[:2], np.linalg.solve(G, q), 1e-9)
        metric = np.diag([np.cos(x[1])**2, 1.])
        check("CP K/U angular metric update", np.trace(metric@(CU[:2, :2]-CK[:2, :2])),
              v[:2]@metric@v[:2]/h, 1e-9)
        if len(selected) == 1:
            check("SS spatial-temporal orthogonality", I[:2, 2:], np.zeros((2, 2)), 1e-10)
            check("SS exact angular subblock equality", CU[:2, :2], CK[:2, :2], 1e-10)
            assert CU[2, 2] > CK[2, 2], "First-sample time origin couples frequency and rate"
    # Nonzero EFIM loss can affect only Doppler after the covariance update.
    K = np.array([[3., .3, .5], [.3, 2., -.4], [.5, -.4, 4.]])
    CK = np.linalg.inv(K)
    for v in [np.array([0., 0., .7]), np.array([.2, -.3, .7])]:
        coupling = K@v
        rate_info = coupling@v+2.
        CU = np.linalg.inv(K-np.outer(coupling, coupling)/rate_info)
        check("synthetic rate-only covariance update", CU-CK, np.outer(v, v)/2., 1e-12)
        assert np.linalg.norm(coupling[:2]) > 0
        if np.linalg.norm(v[:2]) == 0:
            check("nonzero EFIM loss with equal angle bounds", CU[:2, :2], CK[:2, :2], 1e-12)
            assert CU[2, 2] > CK[2, 2]
        else:
            assert np.trace(metric@(CU[:2, :2]-CK[:2, :2])) > 0
    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)
    for ell in range(2):
        p, pd, vu, w, wd, proxy, b, bg, bd, bx = mixed_motion(x, .41, ell)
        step = 1e-3
        check("mixed position velocity", pd,
              (mixed_motion(x, .41+step, ell)[0]-mixed_motion(x, .41-step, ell)[0])/(2*step),
              2e-6)
        check("mixed true LOS time derivative", wd,
              (mixed_motion(x, .41+step, ell)[3]-mixed_motion(x, .41-step, ell)[3])/(2*step))
        check("mixed node Doppler Jacobian", bx,
              finite_difference(lambda x: mixed_motion(x, .41, ell)[6], x))
        check("mixed Doppler time derivative", bd,
              (mixed_motion(x, .41+step, ell)[6]-mixed_motion(x, .41-step, ell)[6])/(2*step))
        check("mixed Doppler discrepancy", b-bg, (vu-pd)@w/WAVELENGTH)
        check("mixed proxy LOS discrepancy", proxy-wd,
              -(np.eye(3)-np.outer(w, w))@(vu-pd)/np.linalg.norm(
                  SATELLITES[ell]+.41*VELOCITIES[ell]+.5*.41**2*ACCELERATIONS[ell]-p))
        assert np.linalg.norm(vu-pd) > 1, "Fixture must expose the mixed protocol"
        e = -w
        check("orientation steering identity",
              np.exp(-2j*np.pi*ARRAY@w/WAVELENGTH),
              np.exp(2j*np.pi*ARRAY@e/WAVELENGTH))
        check("orientation Doppler identity", b, (VELOCITIES[ell]+.41*ACCELERATIONS[ell]-vu)@e/WAVELENGTH)
    mixed_coeff = lambda x: np.linalg.lstsq(DESIGN, [
        mixed_motion(x, t, 1)[6]-mixed_motion(x, t, 0)[6] for t in TIMES], rcond=None)[0]
    mixed_jac = np.linalg.lstsq(DESIGN, [
        mixed_motion(x, t, 1)[9]-mixed_motion(x, t, 0)[9] for t in TIMES], rcond=None)[0]
    check("mixed fixed-design OLS Jacobian", mixed_jac, finite_difference(mixed_coeff, x))
    # Appendix G: independent block phases remove constants, not intra-block slopes.
    for nodes, starts, expected_rank in [([0.], [0., 3.], 0),
                                         ([0., 2.], [0.], 1),
                                         ([0., 2.], [0., 3.], 2),
                                         ([0., 2.], [0., 3e-4], 2)]:
        blocks, weights, centers = [], [], []
        n = np.array(nodes)
        for ell in range(2):
            f, rate = z[2:]+coefficients(x, ell)[0]
            for start in starts:
                w = geometry(x, start, ell)[0]
                a = np.exp(-2j*np.pi*ARRAY@w/WAVELENGTH)
                mu = np.kron(np.exp(2j*np.pi*((f+rate*start)*n+.5*rate*n*n)), a)
                J = np.column_stack((1j*np.repeat(n, len(a))*mu,
                                     1j*np.repeat(start*n+.5*n*n, len(a))*mu))
                projected = J-mu[:, None]*(mu.conj()@J/(mu.conj()@mu))[None, :]
                expected = np.column_stack((
                    1j*np.repeat(n-n.mean(), len(a))*mu,
                    1j*np.repeat(start*(n-n.mean())+.5*(n*n-(n*n).mean()), len(a))*mu))
                check("IP centered temporal scores", projected, expected)
                blocks.append(projected)
                weights.append(2*np.linalg.norm(projected[:, 0])**2)
                centers.append(start+(n[-1]/2))
        J = np.vstack(blocks)
        real_J = np.vstack((J.real, J.imag))*np.sqrt(2)
        assert np.linalg.matrix_rank(real_J, tol=1e-10) == expected_rank
        if len(n) == 2:
            weights, centers = np.array(weights), np.array(centers)
            I = real_J.T@real_J
            vectors = np.column_stack((np.ones(len(centers)), centers))
            check("two-sample temporal information", I, vectors.T@(weights[:, None]*vectors))
            residual = real_J[:, 0]-real_J[:, 1]*(I[1, 0]/I[1, 1])
            check("real rate Schur complement", residual@residual,
                  I[0, 0]-I[0, 1]**2/I[1, 1])
            if expected_rank == 1:
                check("single-frame exact frequency/rate cancellation", J@np.array([1., -1.]),
                      np.zeros(len(J)))
            elif starts[-1] < 1e-3:
                assert 0 < residual@residual/I[0, 0] < 1e-7
    print("PASS: synthetic local formula and temporal-rank checks; code-side geometry and CP/IP acceptance still required.")


if __name__ == "__main__":
    main()