"""Check a conversion against recorded PyTorch FP32 outputs, one model per process.

License: CC-BY-NC-SA-4.0. Checks are numerical smoke tests, not detector accuracy.
"""
import argparse
import hashlib
import json
import platform
import time
from pathlib import Path

import numpy as np
import onnxruntime as ort


def softmax(x):
    ex = np.exp(x.astype(np.float64) - x.max(axis=-1, keepdims=True))
    return ex / ex.sum(axis=-1, keepdims=True)


def digest(path):
    with path.open("rb") as f:
        return hashlib.file_digest(f, "sha256").hexdigest()


def main():
    ort.disable_telemetry_events()
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("variant", choices=["fp32", "fp16", "int8"])
    p.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1])
    args = p.parse_args()
    files = {"fp32": "model.onnx", "fp16": "model_fp16.onnx", "int8": "model_int8.onnx"}
    limits = {"fp32": 0.0001, "fp16": 0.01, "int8": 0.05}
    config = ort.SessionOptions()
    config.intra_op_num_threads = 4
    config.inter_op_num_threads = 1
    config.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
    started = time.perf_counter()
    session = ort.InferenceSession(str(args.root / "onnx" / files[args.variant]),
                                  sess_options=config, providers=["CPUExecutionProvider"])
    load_seconds = time.perf_counter() - started
    reference_info = json.loads((args.root / "validation/reference.json").read_text())
    reference_hash = digest(args.root / "validation/reference.npz")
    fixture_hash = digest(args.root / "validation/fixtures.json")
    if reference_hash != reference_info["reference_npz_sha256"] or fixture_hash != reference_info["fixtures_sha256"]:
        raise RuntimeError("Reference tensors or fixtures changed; regenerate reference outputs.")
    cases = reference_info["cases"]
    reference = np.load(args.root / "validation/reference.npz", allow_pickle=False)
    rows = []
    for case in cases:
        key = case["id"]
        feeds = {name: reference[key + "_" + name] for name in ["input_ids", "attention_mask"]}
        started = time.perf_counter()
        actual = session.run(["logits"], feeds)[0]
        elapsed = time.perf_counter() - started
        expected = reference[key + "_logits"]
        if actual.shape != expected.shape or not np.isfinite(actual).all():
            raise RuntimeError("Invalid output for " + key)
        rows.append({"id": key, "shape": case["shape"],
                     "max_absolute_logit_difference": float(np.max(np.abs(actual - expected))),
                     "max_absolute_probability_difference": float(np.max(np.abs(softmax(actual) - softmax(expected)))),
                     "argmax_agreements": int(np.sum(actual.argmax(-1) == expected.argmax(-1))),
                     "samples": len(actual), "single_run_seconds": elapsed})
        print(args.variant, key, rows[-1], flush=True)
    worst = max(row["max_absolute_probability_difference"] for row in rows)
    disagreements = sum(row["samples"] - row["argmax_agreements"] for row in rows)
    report = {"variant": args.variant, "model": "onnx/" + files[args.variant],
              "model_sha256": digest(args.root / "onnx" / files[args.variant]),
              "reference_npz_sha256": reference_hash, "fixtures_sha256": fixture_hash,
              "provider": "CPUExecutionProvider", "onnxruntime": ort.__version__,
              "platform": {"os": platform.system(), "version": platform.mac_ver()[0],
                           "machine": platform.machine()}, "threads": 4,
              "load_seconds": load_seconds, "cases": rows,
              "max_absolute_probability_difference": worst, "argmax_disagreements": disagreements,
              "samples": sum(row["samples"] for row in rows),
              "acceptance_probability_tolerance": limits[args.variant],
              "passed": worst <= limits[args.variant] and disagreements == 0,
              "limitations": "Synthetic unlabeled conversion fixtures. Timing is one run per shape, not a comparative performance benchmark. No Windows, Linux, CUDA, DirectML, WinML or CoreML validation is implied."}
    (args.root / "validation" / (args.variant + ".json")).write_text(json.dumps(report, indent=2) + "\n")
    if not report["passed"]:
        raise SystemExit("Conversion acceptance check failed; inspect report before publication.")


if __name__ == "__main__":
    main()