"""Correctness gates: re-derive quantities the paper (or classical theory)
already establishes, and assert them before interpreting anything else."""
import json
import numpy as np
import rdpg

OUT = {}


def gate_eig_residual(rng):
    """LAPACK eigendecomposition of A is exact to machine precision."""
    n, r = 800, 5
    A, X = rdpg.weighted_rdpg(n, r, rng)
    s, U = rdpg.full_spectrum(A)
    res = float(np.max(np.abs(A - (U * s) @ U.T)))
    orth = float(np.max(np.abs(U.T @ U - np.eye(n))))
    return {"max_abs_recon_residual": res, "max_abs_orthonormality_residual": orth,
            "n": n, "pass": res < 1e-9 and orth < 1e-9}


def gate_dirichlet_moments(rng):
    """For X_i ~ Dir(1,1,1,1,1), M = E[x x^T] = (I + J)/30 exactly, so the
    non-zero eigenvalues of P = X X^T obey s_1/n -> 1/5, s_2..s_5/n -> 1/30.
    This pins the constant in Assumption A4 (s_j asymp rho_n n) and shows that
    the paper's own simulation design has a 4-fold degenerate population
    eigenvalue (Assumptions A3/A6 hold only through O(sqrt(n)) fluctuations)."""
    r = 5
    M_theory = (np.eye(r) + np.ones((r, r))) / 30.0
    ev_theory = np.sort(np.linalg.eigvalsh(M_theory))[::-1]
    rows = []
    for n in [2000, 20000, 200000]:
        X = rdpg.dirichlet_latent(n, r, rng)
        M_emp = X.T @ X / n
        ev_emp = np.sort(np.linalg.eigvalsh(M_emp))[::-1]
        rows.append({"n": n,
                     "max_abs_moment_err": float(np.max(np.abs(M_emp - M_theory))),
                     "eigs_emp": ev_emp.tolist()})
    return {"M_theory_eigs": ev_theory.tolist(), "rows": rows,
            "degenerate_multiplicity": 4,
            "pass": rows[-1]["max_abs_moment_err"] < 5e-3}


def gate_eckart_young(rng):
    """Lemma 2.1 (k<0 branch) is an Eckart-Young statement.  With E = 0 the
    d-dimensional ASE IS the rank-d truncated eigendecomposition of P, so
    min_W ||Xhat^o W - rho^{1/2}X||_F must EQUAL sqrt(sum_{j>d} s_j) exactly,
    and the (2,inf) lower bound of Lemma 2.1 must hold."""
    n, r, rho = 1500, 5, 1.0
    X = rdpg.dirichlet_latent(n, r, rng)
    P = rho * (X @ X.T)
    s, U = rdpg.full_spectrum(P)
    s_nz = np.sort(np.linalg.eigvalsh(rho * (X.T @ X)))[::-1]      # exact non-zero eigs
    rows = []
    for d in [1, 2, 3, 4]:
        Xh = rdpg.ase_from_spectrum(s, U, d)
        e2i, D = rdpg.err_2inf(Xh, np.sqrt(rho) * X)
        fro = float(np.linalg.norm(D))
        fro_theory = float(np.sqrt(s_nz[d:].sum()))
        bound = float(np.sqrt(s_nz[d:].sum() / n))                  # Lemma 2.1
        rows.append({"d": d, "k": d - r, "frob_measured": fro,
                     "frob_eckart_young": fro_theory,
                     "rel_frob_residual": abs(fro - fro_theory) / fro_theory,
                     "err_2inf": e2i, "lemma21_lower_bound": bound,
                     "ratio_2inf_over_bound": e2i / bound})
    return {"n": n, "rows": rows,
            "max_rel_frob_residual": max(x["rel_frob_residual"] for x in rows),
            "min_ratio": min(x["ratio_2inf_over_bound"] for x in rows),
            "pass": max(x["rel_frob_residual"] for x in rows) < 1e-10
                    and min(x["ratio_2inf_over_bound"] for x in rows) >= 1.0}


def gate_procrustes(rng):
    """W = U V^T from SVD(Xhat^T X) solves min_{W in O} ||Xhat W - X||_F.
    Verified by (a) first-order optimality: W^T Xhat^T X is symmetric PSD,
    (b) beating 20000 random rotations."""
    n, r, d = 400, 5, 7
    A, X = rdpg.weighted_rdpg(n, r, rng)
    s, U = rdpg.full_spectrum(A)
    Xh = rdpg.ase_from_spectrum(s, U, d)
    Xp = np.hstack([X, np.zeros((n, d - r))])
    Uu, sv, Vt = np.linalg.svd(Xh.T @ Xp)
    W = Uu @ Vt
    best = float(np.linalg.norm(Xh @ W - Xp))
    Msym = W.T @ (Xh.T @ Xp)
    sym_err = float(np.max(np.abs(Msym - Msym.T)))
    psd_min = float(np.min(np.linalg.eigvalsh((Msym + Msym.T) / 2)))
    worst = 0.0
    for _ in range(20000):
        Q, _ = np.linalg.qr(rng.standard_normal((d, d)))
        worst = max(worst, best - float(np.linalg.norm(Xh @ Q - Xp)))
    return {"frob_at_procrustes": best, "symmetry_residual": sym_err,
            "min_eig_of_WtXhatTX": psd_min,
            "best_random_rotation_improvement": worst,
            "pass": sym_err < 1e-9 and psd_min > -1e-9 and worst <= 1e-9}


def gate_semicircle_edge(rng):
    """||E|| / (2 sigma sqrt(n)) -> 1 (Wigner edge).  The paper bounds the
    trailing ASE eigenvalues by ||E|| (Sec. 2.1), so this constant matters."""
    rows = []
    for n in [500, 1000, 2000, 4000]:
        vals = []
        for _ in range(3):
            E = rdpg.sym_noise(n, rng, "normal", 1.0)
            vals.append(float(np.linalg.norm(E, 2)))
        rows.append({"n": n, "opnorm": float(np.mean(vals)),
                     "ratio_to_2sqrtn": float(np.mean(vals)) / (2 * np.sqrt(n))})
    return {"rows": rows, "final_ratio": rows[-1]["ratio_to_2sqrtn"],
            "pass": abs(rows[-1]["ratio_to_2sqrtn"] - 1.0) < 0.05}


def gate_signal_incoherence(rng):
    """Corollary B.3.1 / Assumption A1: ||U_{1:r}||_{2,inf} <~ 1/sqrt(n).
    We report the constant sqrt(n)*||U_{1:r}||_{2,inf}, which must not grow."""
    rows = []
    for n in [500, 1000, 2000, 4000, 8000]:
        X = rdpg.dirichlet_latent(n, 5, rng)
        P = X @ X.T
        s, U = rdpg.full_spectrum(P)
        c = rdpg.two_inf(U[:, :5]) * np.sqrt(n)
        rows.append({"n": n, "sqrtn_times_2inf": float(c)})
    return {"rows": rows,
            "drift": rows[-1]["sqrtn_times_2inf"] / rows[0]["sqrtn_times_2inf"],
            "pass": rows[-1]["sqrtn_times_2inf"] < 3 * rows[0]["sqrtn_times_2inf"]}


def jsonable(o):
    if isinstance(o, dict):
        return {k: jsonable(v) for k, v in o.items()}
    if isinstance(o, (list, tuple)):
        return [jsonable(v) for v in o]
    if isinstance(o, (np.bool_, bool)):
        return bool(o)
    if isinstance(o, (np.integer,)):
        return int(o)
    if isinstance(o, (np.floating,)):
        return float(o)
    return o


if __name__ == "__main__":
    rng = np.random.default_rng(20260728)
    OUT["matmul_sanity"] = rdpg.matmul_sanity()
    OUT["eig_residual"] = gate_eig_residual(rng)
    OUT["dirichlet_moments"] = gate_dirichlet_moments(rng)
    OUT["eckart_young"] = gate_eckart_young(rng)
    OUT["procrustes"] = gate_procrustes(rng)
    OUT["semicircle_edge"] = gate_semicircle_edge(rng)
    OUT["signal_incoherence"] = gate_signal_incoherence(rng)
    OUT["all_pass"] = bool(all(v["pass"] for v in OUT.values() if isinstance(v, dict)))
    OUT = jsonable(OUT)
    with open("outputs/gates.json", "w") as f:
        json.dump(OUT, f, indent=1)
    print(json.dumps(OUT, indent=1))