#!/usr/bin/env python3
"""Clean-room exact audit of the three WIRE automatic claims."""

from __future__ import annotations

import hashlib
import json
import platform
import time
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import scipy
from numpy.polynomial.hermite import hermgauss

np.seterr(all="ignore")


HERE = Path(__file__).resolve().parent
OUTPUT_DIR = HERE / "outputs" / "wire"
RESULTS = OUTPUT_DIR / "wire_results.json"
FIGURE = OUTPUT_DIR / "wire_audit.png"
PAPER_SHA = "3eb6899ac0da995483dfba1eafe1ed625a1673d638d498882bcf741709f3415f"
OFFICIAL_COMMIT = "4ac067eb38272543b0cdd7591d630399ff37bce4"


def laplacian_from_edges(n, edges):
    a = np.zeros((n, n), dtype=float)
    for i, j in edges:
        a[i, j] = a[j, i] = 1.0
    return np.diag(a.sum(axis=1)) - a


def path_laplacian(n):
    return laplacian_from_edges(n, [(i, i + 1) for i in range(n - 1)])


def cycle_laplacian(n):
    return laplacian_from_edges(n, [(i, (i + 1) % n) for i in range(n)])


def star_laplacian(n):
    return laplacian_from_edges(n, [(0, i) for i in range(1, n)])


def grid_laplacian(nx, ny):
    edges = []
    for y in range(ny):
        for x in range(nx):
            i = y * nx + x
            if x + 1 < nx:
                edges.append((i, i + 1))
            if y + 1 < ny:
                edges.append((i, i + nx))
    return laplacian_from_edges(nx * ny, edges)


def random_connected_laplacian(n, rng, p=0.25):
    edges = {(i, i + 1) for i in range(n - 1)}
    for i in range(n):
        for j in range(i + 2, n):
            if rng.random() < p:
                edges.add((i, j))
    return laplacian_from_edges(n, sorted(edges))


def graph_library():
    rng = np.random.default_rng(2606)
    out = []
    for n in [6, 9, 13]:
        out.extend([(f"path-{n}", path_laplacian(n)), (f"cycle-{n}", cycle_laplacian(n)), (f"star-{n}", star_laplacian(n))])
    for nx, ny in [(3, 3), (3, 5), (4, 6)]:
        out.append((f"grid-{nx}x{ny}", grid_laplacian(nx, ny)))
    for n in [8, 12, 18]:
        out.append((f"random-{n}", random_connected_laplacian(n, rng)))
    return out


def spectral_coordinates(L, resistance_scaled=False):
    vals, vecs = np.linalg.eigh(L)
    keep = vals > 1e-10
    vals, vecs = vals[keep], vecs[:, keep]
    if resistance_scaled:
        vecs = vecs / np.sqrt(vals)[None, :]
    return vals, vecs


def rotate_rows(x, theta):
    """Apply independent 2D rotations to consecutive feature pairs."""
    pairs = x.reshape(x.shape[0], -1, 2)
    c, s = np.cos(theta), np.sin(theta)
    out = np.empty_like(pairs)
    out[..., 0] = pairs[..., 0] * c - pairs[..., 1] * s
    out[..., 1] = pairs[..., 0] * s + pairs[..., 1] * c
    return out.reshape(x.shape)


def claim1_spectral_structure():
    rows = []
    rng = np.random.default_rng(1807)
    for name, L in graph_library():
        vals, vecs = spectral_coordinates(L)
        m = min(5, vecs.shape[1])
        coords = vecs[:, :m]
        omega = np.linspace(0.7, 1.3, m)
        theta = coords @ omega
        score = np.cos(theta[:, None] - theta[None, :])

        perm = rng.permutation(len(L))
        P = np.eye(len(L))[perm]
        Lp = P @ L @ P.T
        vp = P @ vecs[:, :m]
        residual = np.linalg.norm(Lp @ vp - vp * vals[:m][None, :], ord="fro")
        score_p = np.cos((P @ theta)[:, None] - (P @ theta)[None, :])
        equivariance_error = np.max(np.abs(score_p - P @ score @ P.T))

        q = np.tile([1.0, 0.0], (len(L), 1))
        q_rot = rotate_rows(q, theta[:, None])
        direct_score = q_rot @ q_rot.T
        relative_rotation_error = np.max(np.abs(direct_score - score))
        offdiag = score[~np.eye(len(L), dtype=bool)]
        rows.append(
            {
                "graph": name,
                "nodes": len(L),
                "spectral_dimensions": m,
                "permuted_eigen_residual_fro": float(residual),
                "permutation_equivariance_error": float(equivariance_error),
                "absolute_vs_relative_rotation_error": float(relative_rotation_error),
                "offdiagonal_score_standard_deviation": float(np.std(offdiag)),
                "distinct_scores_rounded_1e-10": int(len(np.unique(np.round(offdiag, 10)))),
            }
        )
    return {
        "protocol": "15 connected graphs; Laplacian eigenfeatures drive pairwise rotary logits; fixed node permutations are checked algebraically",
        "rows": rows,
        "graphs": len(rows),
        "max_permuted_eigen_residual_fro": max(r["permuted_eigen_residual_fro"] for r in rows),
        "max_permutation_equivariance_error": max(r["permutation_equivariance_error"] for r in rows),
        "max_relative_rotation_error": max(r["absolute_vs_relative_rotation_error"] for r in rows),
        "min_offdiagonal_score_standard_deviation": min(r["offdiagonal_score_standard_deviation"] for r in rows),
        "min_distinct_scores": min(r["distinct_scores_rounded_1e-10"] for r in rows),
    }


def claim2_grids_and_resistance():
    path_rows = []
    for n in [4, 5, 8, 13, 21, 34, 55, 64]:
        vals, vecs = np.linalg.eigh(path_laplacian(n))
        analytic = -np.cos(np.pi * (np.arange(n) + 0.5) / n)
        analytic /= np.linalg.norm(analytic)
        corr = abs(float(analytic @ vecs[:, 1]))
        eigenvalue_error = abs(vals[1] - (2.0 - 2.0 * np.cos(np.pi / n)))
        path_rows.append({"n": n, "absolute_mode_correlation": corr, "eigenvalue_error": float(eigenvalue_error)})

    grid_rows = []
    for nx, ny in [(3, 4), (3, 7), (4, 4), (4, 9), (5, 6), (6, 8), (7, 9), (8, 10)]:
        L = grid_laplacian(nx, ny)
        xx = np.tile(np.arange(nx), ny)
        yy = np.repeat(np.arange(ny), nx)
        ux = np.cos(np.pi * (xx + 0.5) / nx)
        uy = np.cos(np.pi * (yy + 0.5) / ny)
        lamx = 2.0 - 2.0 * np.cos(np.pi / nx)
        lamy = 2.0 - 2.0 * np.cos(np.pi / ny)
        rx = np.linalg.norm(L @ ux - lamx * ux) / np.linalg.norm(ux)
        ry = np.linalg.norm(L @ uy - lamy * uy) / np.linalg.norm(uy)
        grid_rows.append({"grid": f"{nx}x{ny}", "x_mode_residual": float(rx), "y_mode_residual": float(ry)})

    hx, hw = hermgauss(80)
    resistance_rows, slopes = [], []
    omega_grid = np.asarray([0.01, 0.02, 0.04, 0.08, 0.16])
    for name, L in graph_library():
        vals, coords = spectral_coordinates(L, resistance_scaled=True)
        pinv = np.linalg.pinv(L, hermitian=True)
        diag = np.diag(pinv)
        R_pinv = diag[:, None] + diag[None, :] - 2.0 * pinv
        R_coords = np.sum((coords[:, None, :] - coords[None, :, :]) ** 2, axis=2)
        resistance_error = np.max(np.abs(R_pinv - R_coords))
        pairs = [(i, j) for i in range(len(L)) for j in range(i + 1, len(L))]
        # Cover low, median, and high resistance pairs deterministically.
        pairs.sort(key=lambda ij: R_pinv[ij])
        selected = [pairs[0], pairs[len(pairs) // 2], pairs[-1]]
        for i, j in selected:
            R = float(R_pinv[i, j])
            exact = np.exp(-0.5 * omega_grid**2 * R)
            leading = 1.0 - 0.5 * omega_grid**2 * R
            residuals = np.abs(exact - leading)
            slope = float(np.polyfit(np.log(omega_grid), np.log(residuals), 1)[0])
            slopes.append(slope)
            # Independent Gauss-Hermite expectation E cos(omega*sqrt(R)*Z).
            gh = np.array([np.sum(hw * np.cos(w * np.sqrt(R) * np.sqrt(2) * hx)) / np.sqrt(np.pi) for w in omega_grid])
            resistance_rows.append(
                {
                    "graph": name,
                    "pair": [i, j],
                    "effective_resistance": R,
                    "spectral_vs_pseudoinverse_max_error_graph": float(resistance_error),
                    "small_omega_remainder_loglog_slope": slope,
                    "max_gauss_hermite_vs_exact_expectation_error": float(np.max(np.abs(gh - exact))),
                    "remainder_over_omega4_at_smallest_omega": float(residuals[0] / omega_grid[0] ** 4),
                }
            )
    return {
        "path_grid_protocol": "closed-form first path/grid Laplacian modes, up to sign and normalization",
        "path_rows": path_rows,
        "grid_rows": grid_rows,
        "min_path_mode_correlation": min(r["absolute_mode_correlation"] for r in path_rows),
        "max_path_eigenvalue_error": max(r["eigenvalue_error"] for r in path_rows),
        "max_grid_mode_residual": max(max(r["x_mode_residual"], r["y_mode_residual"]) for r in grid_rows),
        "resistance_protocol": "Laplacian-pseudoinverse resistance vs squared resistance-scaled spectral distance; exact Gaussian WIRE expectation and independent 80-node Gauss-Hermite oracle",
        "resistance_rows": resistance_rows,
        "pairs": len(resistance_rows),
        "max_spectral_vs_pseudoinverse_resistance_error": max(r["spectral_vs_pseudoinverse_max_error_graph"] for r in resistance_rows),
        "remainder_slope_range": [min(slopes), max(slopes)],
        "max_gauss_hermite_expectation_error": max(r["max_gauss_hermite_vs_exact_expectation_error"] for r in resistance_rows),
    }


def claim3_linear_attention():
    rng = np.random.default_rng(314159)
    rows = []
    for n in [8, 32, 128, 512, 1024, 2048]:
        d, dv, m = 32, 7, 6
        q, k, v = rng.normal(size=(n, d)), rng.normal(size=(n, d)), rng.normal(size=(n, dv))
        coords, omega = rng.normal(size=(n, m)), rng.normal(size=(d // 2, m))
        theta = coords @ omega.T
        qr, kr = rotate_rows(q, theta), rotate_rows(k, theta)
        associative = qr @ (kr.T @ v)
        pairwise = (qr @ kr.T) @ v
        relerr = np.linalg.norm(associative - pairwise) / np.linalg.norm(pairwise)
        dense_elements = n * n
        associative_elements = d * dv
        rows.append(
            {
                "tokens": n,
                "feature_dim": d,
                "value_dim": dv,
                "relative_associativity_error": float(relerr),
                "dense_attention_elements": dense_elements,
                "associative_intermediate_elements": associative_elements,
                "intermediate_memory_ratio": dense_elements / associative_elements,
                "dense_leading_multiply_adds": int(2 * n * n * d),
                "associative_leading_multiply_adds": int(2 * n * d * dv),
            }
        )
    return {
        "identity": "(Q_rot K_rot^T)V = Q_rot(K_rot^T V); WIRE acts on Q and K before the associative linear-attention contraction",
        "rows": rows,
        "max_relative_associativity_error": max(r["relative_associativity_error"] for r in rows),
        "largest_n_memory_ratio": rows[-1]["intermediate_memory_ratio"],
        "largest_n_operation_ratio": rows[-1]["dense_leading_multiply_adds"] / rows[-1]["associative_leading_multiply_adds"],
    }


def make_figure(c2, c3):
    fig, axes = plt.subplots(1, 3, figsize=(15, 4.4))
    ns = [r["n"] for r in c2["path_rows"]]
    corr_err = [1 - r["absolute_mode_correlation"] for r in c2["path_rows"]]
    axes[0].semilogy(ns, np.maximum(corr_err, 1e-17), "o-")
    axes[0].set(xlabel="path nodes", ylabel="1 − |mode correlation|", title="Grid modes recover RoPE coordinates")
    axes[0].grid(alpha=0.25)

    slopes = [r["small_omega_remainder_loglog_slope"] for r in c2["resistance_rows"]]
    axes[1].hist(slopes, bins=10, color="#6d5dfc", alpha=0.85)
    axes[1].axvline(4, color="black", linestyle="--", label="O(ω⁴)")
    axes[1].set(xlabel="fitted remainder exponent", ylabel="graph-pair count", title="Effective-resistance expansion")
    axes[1].legend()

    nr = [r["tokens"] for r in c3["rows"]]
    mr = [r["intermediate_memory_ratio"] for r in c3["rows"]]
    axes[2].loglog(nr, mr, "o-", color="#e45756")
    axes[2].set(xlabel="tokens N", ylabel="dense / associative intermediate", title="WIRE keeps linear attention associative")
    axes[2].grid(alpha=0.25, which="both")
    fig.tight_layout()
    fig.savefig(FIGURE, dpi=180)


def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    t0 = time.time()
    c1 = claim1_spectral_structure()
    c2 = claim2_grids_and_resistance()
    c3 = claim3_linear_attention()

    assertions = {
        "claim1_all_graphs_structurally_nonconstant": c1["min_offdiagonal_score_standard_deviation"] > 1e-5 and c1["min_distinct_scores"] >= 3,
        "claim1_equivariance_and_rotation_identities": c1["max_permutation_equivariance_error"] < 1e-12 and c1["max_relative_rotation_error"] < 1e-12,
        "claim2_path_and_grid_modes": c2["min_path_mode_correlation"] > 1 - 1e-12 and c2["max_grid_mode_residual"] < 1e-12,
        "claim2_resistance_identity": c2["max_spectral_vs_pseudoinverse_resistance_error"] < 1e-11,
        "claim2_omega4_remainder": c2["remainder_slope_range"][0] > 3.94 and c2["remainder_slope_range"][1] < 4.01,
        "claim2_expectation_oracle": c2["max_gauss_hermite_expectation_error"] < 1e-12,
        "claim3_linear_attention_associativity": c3["max_relative_associativity_error"] < 1e-11,
        "claim3_subquadratic_intermediate": c3["largest_n_memory_ratio"] > 10000,
    }
    assert all(assertions.values()), assertions
    make_figure(c2, c3)

    payload = {
        "paper": {
            "title": "Rotary Position Encodings for Graphs",
            "openreview": "trn64znfNx",
            "arxiv": "2509.22259",
            "paper_pdf_sha256": PAPER_SHA,
            "official_repository": "https://github.com/cederikhoefs/Graph-RoPE",
            "official_commit": OFFICIAL_COMMIT,
            "relationship": "clean-room NumPy/SciPy implementation of paper equations; author code not executed",
        },
        "automatic_claims": {
            "claim_1_spectral_structural_rotation": c1,
            "claim_2_grid_rope_and_effective_resistance": c2,
            "claim_3_linear_attention_compatibility": c3,
        },
        "assertions": assertions,
        "limitations": [
            "Claims are audited at the mathematical operator level; this does not retrain the paper's more than 200 graph models.",
            "The effective-resistance statement is the paper's random-frequency small-omega expectation, not exact invariance for one learned frequency draw.",
            "The linear-attention check targets the bilinear numerator factorization; normalization can be computed separately as in standard linear attention.",
        ],
        "environment": {"python": platform.python_version(), "numpy": np.__version__, "scipy": scipy.__version__},
        "runtime_seconds": time.time() - t0,
    }
    RESULTS.write_text(json.dumps(payload, indent=2) + "\n")
    print(json.dumps({"assertions": assertions, "runtime_seconds": payload["runtime_seconds"]}, indent=2))


if __name__ == "__main__":
    main()