import json
import numpy as np


def spectral_gate(A, c, tau=1e-3):
    c = np.asarray(c, float); A = np.asarray(A, float)
    lam = float(c @ A @ c / (c @ c))
    residual = float(np.linalg.norm(c @ A - lam * c) / np.linalg.norm(c))
    return residual, lam, (residual > tau or lam < 0)


def powers_and_cov(A, P0, Q, H):
    """Return propagated covariances P_j, including P_0."""
    P = np.asarray(P0, float).copy(); out = [P.copy()]
    for _ in range(H):
        P = A @ P @ A.T + Q
        out.append(P.copy())
    return out


def verify_math(rng):
    d = 4
    # An exact left eigenvector, and a non-eigenvector rotating readout.
    A = np.diag([0.82, 0.61, 0.45, 0.30])
    c = np.array([1., 0., 0., 0.])
    lam = .82
    max_power_err = 0.
    for j in range(7):
        max_power_err = max(max_power_err, np.linalg.norm(c @ np.linalg.matrix_power(A, j) - lam**j*c))
    r, l, skip = spectral_gate(A, c)
    # Covariance recursion versus the closed finite sum.
    P0 = np.diag([.2, .1, .15, .05]); Q = .01*np.eye(d); H = 6
    rec = powers_and_cov(A, P0, Q, H)
    closed = []
    for j in range(H+1):
        Aj = np.linalg.matrix_power(A, j)
        Pj = Aj @ P0 @ Aj.T
        for i in range(j):
            Ai = np.linalg.matrix_power(A, i)
            Pj += Ai @ Q @ Ai.T
        closed.append(Pj)
    cov_err = max(np.linalg.norm(rec[j]-closed[j]) for j in range(H+1))
    # A rotation must have a clearly nonzero residual and activate the gate.
    th = .55; R = .98*np.array([[np.cos(th), -np.sin(th)], [np.sin(th), np.cos(th)]])
    rr, ll, active = spectral_gate(R, np.array([1., 0.]), tau=1e-3)
    return dict(power_identity_error=max_power_err, covariance_identity_error=cov_err,
                aligned_residual=r, aligned_lambda=l, aligned_skips=not skip,
                rotating_residual=rr, rotating_lambda=ll, rotating_activates=active)


def rollout_decision(A, x, c, delta, H):
    """Alarm if current or any of H predicted readouts crosses delta."""
    y = np.asarray(x, float).copy(); best = float(c @ y)
    for _ in range(H):
        y = A @ y
        best = max(best, float(c @ y))
    return best >= delta


def benchmark(rng, n=3000, H=8, delta=.8, tau=0.02):
    th = .55
    rotation = .98*np.array([[np.cos(th), -np.sin(th)], [np.sin(th), np.cos(th)]])
    aligned = np.diag([.82, .60])
    c = np.array([1., 0.])
    systems = [("aligned_positive", aligned), ("rotating", rotation)]
    rows = []
    for name, A in systems:
        # Random current states. Deterministic dynamics make the structural result visible.
        X = rng.normal(size=(n, 2))
        truth = np.array([rollout_decision(A, x, c, delta, H) for x in X])
        current = (X @ c >= delta)
        always = np.array([rollout_decision(A, x, c, delta, H) for x in X])
        residual, lam, activates = spectral_gate(A, c, tau)
        gated = current if not activates else always
        # Cost is number of matrix-vector rollout steps, with current readout free in this comparison.
        gate_steps = (H if activates else 0) * n
        always_steps = H*n
        def scores(pred):
            tp = np.sum(pred & truth); fp = np.sum(pred & ~truth); fn = np.sum(~pred & truth)
            return dict(recall=float(tp/max(1,tp+fn)), false_positive_rate=float(fp/max(1,np.sum(~truth))))
        rows.append(dict(system=name, residual=residual, lambda_hat=lam,
                         gate_activates=activates, truth_rate=float(truth.mean()),
                         current=scores(current), always_rollout=scores(always), gated=scores(gated),
                         always_rollout_steps=always_steps, gated_rollout_steps=gate_steps,
                         rollout_step_reduction=float(1-gate_steps/always_steps)))
    # Mixed workload: equal systems, showing selective compute and overall recall.
    return rows


def main():
    rng = np.random.default_rng(12345)
    math = verify_math(rng)
    rows = benchmark(rng)
    result = {"math_verification": math, "benchmark": rows}
    with open("results.json", "w") as f: json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()