# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Write assets/nn_reference.npz: the self-test's outputs in the reference environment, for the near-identical check.

  source env/canonical.sh && python -s tools/make_nn_reference.py --device cuda:0

Runs the start-up self-test's cases (nn/torch_backend.py, golden_cases) in the exact lane and keeps their outputs only
if every one's digest equals the release's (assets/nn_golden.json), so the arrays are the release's own outputs.
Outside the reference environment, a Diarizer that is not strict compares its self-test outputs with these
(torch_backend.NEAR). Stores segment_b8 (the smaller segmentation batches are its first rows) and the voiceprint cases.
"""
import argparse
import hashlib
import io
import json
import sys
from pathlib import Path

import numpy as np

PROJ = Path(__file__).resolve().parents[1]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("--out", type=Path, default=PROJ / "assets/nn_reference.npz")
    a = ap.parse_args()
    sys.path.insert(0, str(PROJ / "python"))
    from anuvani_diar.nn.torch_backend import TorchNN, _digest, golden_cases
    golden = PROJ / "assets/nn_golden.json"
    nn = TorchNN(PROJ / "artifacts/bundle", devices=[a.device], golden=golden, strict=True)
    want = json.loads(golden.read_text())["digests"]
    nn._force_exact = True
    out = {name: fn() for name, fn in golden_cases(nn, 0)}
    bad = [k for k in want if _digest(out[k]) != want[k]]
    if bad:
        raise SystemExit(f"not the release's outputs: {bad}")
    keep = {k: np.ascontiguousarray(v, dtype=np.float32) for k, v in out.items() if not k.startswith("segment") or k == "segment_b8"}
    buf = io.BytesIO()
    np.savez_compressed(buf, **keep)
    a.out.write_bytes(buf.getvalue())
    print(json.dumps(dict(out=str(a.out), bytes=a.out.stat().st_size, sha256=hashlib.sha256(buf.getvalue()).hexdigest(),
                          arrays={k: list(v.shape) for k, v in keep.items()}, device=nn.info()["gpus"][0]), indent=1))


if __name__ == "__main__":
    main()