"""Independent raw-array recomputation for the Claim 5 jaw experiment."""

from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
from scipy.ndimage import gaussian_filter


def check(record: dict) -> dict:
    failures: list[str] = []
    indices = np.asarray(record["voxel_indices"], dtype=np.int64)
    weights = np.asarray(record["weights"], dtype=np.float64)
    scaling = np.asarray(record["scaling"], dtype=np.float64)
    kernel = record["kernel"]
    grid_shape = tuple(record["voxelization"]["grid_shape"])
    workspace = np.zeros(grid_shape, dtype=np.float64)

    def kernel_matvec(vector: np.ndarray) -> np.ndarray:
        workspace.fill(0.0)
        workspace[indices[:, 0], indices[:, 1], indices[:, 2]] = vector
        convolved = gaussian_filter(
            workspace,
            sigma=float(kernel["sigma_grid_cells"]),
            mode="constant",
            cval=0.0,
            truncate=float(kernel["truncate_sigma"]),
        )
        return convolved[
            indices[:, 0], indices[:, 1], indices[:, 2]
        ]

    def apply(signal: np.ndarray) -> np.ndarray:
        return scaling * kernel_matvec(weights * scaling * signal)

    row_values = apply(np.ones(weights.shape[0], dtype=np.float64))
    row_residual = float(np.max(np.abs(row_values - 1.0)))
    if row_residual > 2e-10:
        failures.append("independent constant preservation")
    if abs(float(weights.sum()) - 1.0) > 1e-14:
        failures.append("independent mass weights")
    if np.any(weights <= 0.0) or np.any(scaling <= 0.0):
        failures.append("independent positive weights/scaling")

    source_index = int(record["diffusion"]["source_index"])
    signal = np.zeros(weights.shape[0], dtype=np.float64)
    signal[source_index] = 1.0 / weights[source_index]
    snapshots = {
        int(item["step"]): item
        for item in record["diffusion"]["snapshots"]
    }
    mass_errors: list[float] = []
    signal_errors: list[float] = []
    l2_values: list[float] = []
    roughness_values: list[float] = []
    maximum_step = max(snapshots)
    for step in range(maximum_step + 1):
        next_signal = apply(signal)
        if step in snapshots:
            stored = np.asarray(snapshots[step]["signal"], dtype=np.float64)
            signal_errors.append(
                float(np.max(np.abs(stored - signal)))
            )
            mass = float(np.sum(weights * signal))
            mass_errors.append(abs(mass - 1.0))
            centered = signal - mass
            l2_values.append(float(np.sum(weights * centered**2)))
            roughness_values.append(
                float(np.sum(weights * signal * (signal - next_signal)))
            )
            if float(signal.min()) < -2e-10:
                failures.append(f"independent positivity step {step}")
        if step < maximum_step:
            signal = next_signal

    if max(signal_errors) > 2e-11:
        failures.append("independent stored-signal recomputation")
    if max(mass_errors) > 2e-10:
        failures.append("independent mass conservation")
    if any(
        later > earlier * (1.0 + 2e-9) + 1e-10
        for earlier, later in zip(l2_values, l2_values[1:])
    ):
        failures.append("independent monotone L2 smoothing")
    if min(roughness_values) < -2e-8:
        failures.append("independent nonnegative roughness")
    if any(
        later > earlier * (1.0 + 2e-8) + 1e-10
        for earlier, later in zip(
            roughness_values, roughness_values[1:]
        )
    ):
        failures.append("independent monotone roughness")

    modalities = record["cross_modalities"]
    if modalities["point_cloud"]["count"] != 5_000:
        failures.append("independent point modality")
    gmm = modalities["covariance_aware_gmm"]
    if (
        gmm["count"] != 500
        or gmm.get("covariance_matrices") != 500
        or gmm.get("minimum_covariance_eigenvalue", 0.0) <= 0.0
    ):
        failures.append("independent covariance-aware GMM modality")
    if modalities["sparse_armadillo_voxels"]["nonempty_voxels"] < 100:
        failures.append("independent voxel modality")

    return {
        "claim_5_contract_pass": not failures,
        "verifier": "independent_raw_recomputation",
        "failures": failures,
        "recomputed": {
            "row_residual_max": row_residual,
            "maximum_mass_error": max(mass_errors),
            "maximum_stored_signal_error": max(signal_errors),
            "l2_values": l2_values,
            "roughness_values": roughness_values,
        },
    }


def main() -> None:
    if len(sys.argv) != 3:
        raise SystemExit(
            "usage: check_claim5_independent.py RAW_JSON OUTPUT_JSON"
        )
    record = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8"))
    result = check(record)
    Path(sys.argv[2]).write_text(
        json.dumps(result, indent=2) + "\n", encoding="utf-8"
    )
    print("CLAIM5_INDEPENDENT=" + json.dumps(result, sort_keys=True))
    if not result["claim_5_contract_pass"]:
        raise SystemExit(1)


if __name__ == "__main__":
    main()