# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Write assets/nn_golden.json: the release's own network outputs (sha256 of the float32 bytes) for the self-test
inputs of python/anuvani_diar/nn/torch_backend.py, on every batch shape the engine uses.

The outputs come from the release's code path only: `load_bundle` (frozen package), then
`pipeline._segmentation.infer` (segmentation + the softmax conversion) and `pipeline._embedding(...)` (voiceprints).
The repacked-voiceprint case is checked against the release's two plain voiceprint batches it replaces.

  source env/canonical.sh && CUDA_VISIBLE_DEVICES=0 python -s tools/make_nn_goldens.py
"""
import hashlib
import json
import os
import sys
from pathlib import Path

import numpy as np

PROJ = Path(__file__).resolve().parents[1]
REL = PROJ.parent / "anuvani_diar_v1"
sys.path.insert(0, str(PROJ / "tests"))
sys.path.insert(0, str(PROJ / "python"))


def digest(a):
    return hashlib.sha256(np.ascontiguousarray(a).tobytes()).hexdigest()


def main():
    if os.environ.get("CUBLAS_WORKSPACE_CONFIG") != ":4096:8" or not sys.flags.no_user_site:
        raise SystemExit("source env/canonical.sh and run with python -s")
    import refpkg
    refpkg.load()
    import torch
    from anuvani_diar_ref.frontend import load_bundle
    from anuvani_diar.nn.torch_backend import golden_inputs
    model = load_bundle(REL / "bundle", "cuda:0", 8, 2)
    pipe = model["pipeline"]
    win, masks = golden_inputs()
    d = {}
    with torch.inference_mode():
        for b in range(1, 9):
            d[f"segment_b{b}"] = digest(pipe._segmentation.infer(torch.from_numpy(win[:b].copy())))
        pair = np.concatenate([np.repeat(win[0:1], 4, 0), np.repeat(win[1:2], 4, 0)])
        w8 = np.concatenate([masks[0], masks[1]])
        e8 = pipe._embedding(torch.from_numpy(pair), masks=torch.from_numpy(w8))
        d["embed_b8"] = digest(e8)
        d["embed_b4"] = digest(pipe._embedding(torch.from_numpy(np.repeat(win[2:3], 4, 0)), masks=torch.from_numpy(masks[2].copy())))
        pair2 = np.concatenate([np.repeat(win[2:3], 4, 0), np.repeat(win[3:4], 4, 0)])
        e8b = pipe._embedding(torch.from_numpy(pair2), masks=torch.from_numpy(np.concatenate([masks[2], masks[3]])))
        d["embed_repacked_h2"] = digest(np.concatenate([e8, e8b]))
    out = dict(source="the reference program's networks (reference/ in the model repository), FP32, deterministic",
               torch=torch.__version__, cudnn=torch.backends.cudnn.version(), digests=d)
    (PROJ / "assets/nn_golden.json").write_text(json.dumps(out, indent=1) + "\n")
    print(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()