#!/usr/bin/env python3
"""Independent NumPy oracle for the saved source-scale Deep-UFM state."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
from pathlib import Path

import numpy as np

K = 3
D = 65
N_PER_CLASS = 40
N = K * N_PER_CLASS
LAYER = 4


def digest(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def relu(value: np.ndarray) -> np.ndarray:
    return np.maximum(value, 0.0)


def analyse(state_path: Path) -> dict:
    state = np.load(state_path)
    h = state["H1"].astype(np.float64)
    target = state["Y"].astype(np.float64)
    weights = [
        state[f"W{index}"].astype(np.float64) for index in range(1, 6)
    ]

    # Accelerate-backed NumPy on macOS can leave floating-point status flags
    # set after a finite BLAS matmul and consequently emit spurious divide or
    # overflow RuntimeWarnings on the next operation.  Suppress only those
    # status-flag reports and fail explicitly on every non-finite array.
    with np.errstate(divide="ignore", over="ignore", invalid="ignore"):
        activations = [h]
        preactivations = []
        x = h
        for weight in weights[:-1]:
            z = weight @ x
            preactivations.append(z)
            x = relu(z)
            activations.append(x)
        output = weights[-1] @ x
        residual = output - target
    if not all(
        np.isfinite(value).all()
        for value in [*activations, *preactivations, output, residual]
    ):
        raise RuntimeError("non-finite value in independent forward oracle")

    # For output k and sample j:
    # d output[k,j] / d W4[a,b]
    # = W5[k,a] 1[z4[a,j]>0] activation3[b,j].
    left = weights[-1][:, :, None] * (
        preactivations[3] > 0
    )[None, :, :]
    jacobian = np.einsum(
        "kaj,bj->jkab", left, activations[3], optimize=True
    ).reshape(N * K, D * D)

    # The non-zero eigenvalues of J.T J / N equal those of J J.T / N.
    with np.errstate(divide="ignore", over="ignore", invalid="ignore"):
        gram = (jacobian @ jacobian.T) / N
    if not np.isfinite(gram).all():
        raise RuntimeError("non-finite value in independent Hessian oracle")
    eigenvalues, left_eigenvectors = np.linalg.eigh(
        0.5 * (gram + gram.T)
    )
    order = np.argsort(eigenvalues)[::-1]
    eigenvalues = eigenvalues[order]
    left_eigenvectors = left_eigenvectors[:, order]

    with np.errstate(divide="ignore", over="ignore", invalid="ignore"):
        gradient = (jacobian.T @ residual.T.reshape(-1)) / N
    if not np.isfinite(gradient).all():
        raise RuntimeError("non-finite value in independent gradient oracle")
    gradient_norm = np.linalg.norm(gradient)
    coefficients = []
    for index, eigenvalue in enumerate(eigenvalues[: K * K]):
        with np.errstate(divide="ignore", over="ignore", invalid="ignore"):
            right = (
                jacobian.T @ left_eigenvectors[:, index]
            ) / np.sqrt(max(N * eigenvalue, 1e-300))
        right /= max(np.linalg.norm(right), 1e-300)
        with np.errstate(divide="ignore", over="ignore", invalid="ignore"):
            coefficient = float(
                (right @ gradient) ** 2 / max(gradient_norm**2, 1e-300)
            )
        if not np.isfinite(right).all() or not np.isfinite(coefficient):
            raise RuntimeError("non-finite value in eigenspace oracle")
        coefficients.append(coefficient)

    positive = eigenvalues[eigenvalues > max(eigenvalues[0] * 1e-10, 1e-14)]
    top9 = eigenvalues[: K * K]
    ninth_tenth_ratio = float(top9[-1] / max(eigenvalues[K * K], 1e-300))
    top9_unequal_ratio = float(top9[0] / max(top9[-1], 1e-300))
    coeff_sorted = sorted(coefficients, reverse=True)
    coefficient_threshold = max(coeff_sorted[0] * 1e-4, 1e-12)
    nonzero_coefficients = sum(
        coefficient > coefficient_threshold for coefficient in coefficients
    )

    return {
        "state_sha256": digest(state_path),
        "configuration": {
            "K": K,
            "d": D,
            "n_per_class": N_PER_CLASS,
            "L": 5,
            "audited_layer_l": LAYER,
            "parameter_count_W4": D * D,
            "jacobian_shape": list(jacobian.shape),
            "gram_shape": list(gram.shape),
        },
        "fit": {
            "mse": float(np.mean(residual**2)),
            "accuracy": float(
                np.mean(np.argmax(output, axis=0) == np.argmax(target, axis=0))
            ),
        },
        "hessian": {
            "construction": "independent NumPy J J^T / N exact non-zero spectrum",
            "strictly_positive_eigenvalues": int(len(positive)),
            "top_12_eigenvalues": eigenvalues[:12].tolist(),
            "top_9_eigenvalues": top9.tolist(),
            "ninth_to_tenth_separation_ratio": ninth_tenth_ratio,
            "top9_max_to_min_ratio": top9_unequal_ratio,
            "nine_outlier_gate": ninth_tenth_ratio >= 3.0,
            "unequal_top9_gate": top9_unequal_ratio >= 1.05,
        },
        "gradient": {
            "construction": (
                "non-regularization layer-W4 gradient projected onto the "
                "true top-nine Hessian eigenvectors"
            ),
            "squared_alignment_coefficients": coefficients,
            "sorted_squared_alignment_coefficients": coeff_sorted,
            "nonzero_threshold": coefficient_threshold,
            "nonzero_coefficient_count": nonzero_coefficients,
            "top3_max_to_min_ratio": float(
                coeff_sorted[0] / max(coeff_sorted[2], 1e-300)
            ),
            "fourth_to_third_ratio": float(
                coeff_sorted[3] / max(coeff_sorted[2], 1e-300)
            ),
            "K_nonzero_gate": nonzero_coefficients == K,
            "unequal_top3_gate": (
                coeff_sorted[0] / max(coeff_sorted[2], 1e-300)
            ) >= 1.05,
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--state", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    result = analyse(args.state)
    predicates = {
        "nine_hessian_outliers_separate": result["hessian"][
            "nine_outlier_gate"
        ],
        "nine_hessian_outliers_remain_unequal": result["hessian"][
            "unequal_top9_gate"
        ],
        "gradient_has_exactly_K_nonzero_coefficients": result["gradient"][
            "K_nonzero_gate"
        ],
        "top_K_gradient_coefficients_are_unequal": result["gradient"][
            "unequal_top3_gate"
        ],
    }
    result["literal_claim_predicates"] = predicates
    result["all_literal_claim_gates_pass"] = bool(all(predicates.values()))
    result["falsified_literal_predicates"] = [
        name for name, passed in predicates.items() if not passed
    ]
    native_fit_gate = bool(
        result["fit"]["accuracy"] >= 0.99 and result["fit"]["mse"] <= 1e-4
    )
    result["native_fit_gate"] = native_fit_gate
    if not native_fit_gate:
        result["decisive_literal_verdict"] = "inconclusive"
        result["release_quality_gate_pass"] = False
    elif result["all_literal_claim_gates_pass"]:
        result["decisive_literal_verdict"] = "verified"
        result["release_quality_gate_pass"] = True
    else:
        result["decisive_literal_verdict"] = (
            "falsified_as_literally_registered"
        )
        result["release_quality_gate_pass"] = True
    (args.output / "native_oracle.json").write_text(
        json.dumps(result, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    with (args.output / "native_spectrum.csv").open(
        "w", newline="", encoding="utf-8"
    ) as handle:
        writer = csv.writer(handle)
        writer.writerow(("rank", "hessian_eigenvalue", "gradient_alignment"))
        eigenvalues = result["hessian"]["top_12_eigenvalues"]
        coefficients = result["gradient"]["squared_alignment_coefficients"]
        for index, value in enumerate(eigenvalues, 1):
            writer.writerow(
                (index, value, coefficients[index - 1] if index <= 9 else "")
            )
    print(json.dumps(result, indent=2, sort_keys=True))
    if not result["release_quality_gate_pass"]:
        raise SystemExit(2)


if __name__ == "__main__":
    main()