#!/usr/bin/env python3
"""Normalize paired native-run artifacts without changing model arrays.

The first paired run used NumPy's timestamped ``savez`` container and included
wall-clock/process metadata.  This utility rewrites exactly the saved arrays
with fixed ZIP metadata, removes only nondeterministic metadata, corrects the
immutable source identifier, and then requires the two complete run
directories to be byte-identical.
"""

from __future__ import annotations

import argparse
import hashlib
import io
import json
import zipfile
from pathlib import Path

import numpy as np


ARRAY_ORDER = ("H1", "Y", "W1", "W2", "W3", "W4", "W5")


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


def deterministic_npz(path: Path, arrays: dict[str, np.ndarray]) -> None:
    with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_STORED) as archive:
        for name in ARRAY_ORDER:
            payload = io.BytesIO()
            np.lib.format.write_array(
                payload, np.asanyarray(arrays[name]), allow_pickle=False
            )
            info = zipfile.ZipInfo(f"{name}.npy", (1980, 1, 1, 0, 0, 0))
            info.compress_type = zipfile.ZIP_STORED
            info.external_attr = 0o600 << 16
            archive.writestr(info, payload.getvalue())


def normalize(directory: Path) -> dict[str, str]:
    state_path = directory / "final_state.npz"
    result_path = directory / "training_results.json"
    with np.load(state_path) as loaded:
        if set(loaded.files) != set(ARRAY_ORDER):
            raise RuntimeError(f"unexpected state keys in {state_path}")
        arrays = {name: loaded[name].copy() for name in ARRAY_ORDER}
    if not all(np.isfinite(array).all() for array in arrays.values()):
        raise RuntimeError(f"non-finite state in {state_path}")
    deterministic_npz(state_path, arrays)

    result = json.loads(result_path.read_text(encoding="utf-8"))
    result.pop("runtime_seconds", None)
    implementation = result["implementation"]
    implementation.pop("pid", None)
    result["paper"]["title"] = "Unifying Low Dimensional Spectra in Deep Learning"
    result["paper"]["source_revision"] = "arXiv:2404.06106v1"
    result["final_state_sha256"] = digest(state_path)
    result_path.write_text(
        json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    return {
        "final_state.npz": digest(state_path),
        "training_results.json": digest(result_path),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("run_a", type=Path)
    parser.add_argument("run_b", type=Path)
    args = parser.parse_args()
    hashes_a = normalize(args.run_a)
    hashes_b = normalize(args.run_b)
    if hashes_a != hashes_b:
        raise RuntimeError(
            f"paired replay is not byte-identical: A={hashes_a}, B={hashes_b}"
        )
    print(
        json.dumps(
            {"status": "PASS", "paired_byte_identical": True, "sha256": hashes_a},
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()