# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""check-exact: the byte-identity checks on cached front-end records and on audio inputs, on a stratified set, in about
a minute.

  source env/canonical.sh && python -s tools/check_exact.py [--devices cuda:0,cuda:1] [--sample 20]

1. Every GPU's networks against the release's own outputs on every batch shape (the start-up self-test,
   assets/nn_golden.json).
2. Through `Diarizer` (decode, windows, both networks, masks, records, clustering, decoder, RTTM): the 12 golden files,
   every robustness input (formats, rates, channels, cut lengths, silence, NaN, empty), the 1 h file and a stratified
   sample of benchmark files, against the raw-audio oracle (the release's own code on the same audio): every front-end
   record, every RTTM and every error byte-identical (tests/test_identity_audio.py's comparison).
3. 40 cached front-end records through the Rust post-processing against the cache oracle
   (tests/test_identity_cached.py).

Exit 0 only if every compared byte is identical. Prints one line per part and the time it took.
"""
import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np

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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--devices", default=None, help="comma-separated torch devices (default: every visible GPU)")
    ap.add_argument("--sample", type=int, default=20, help="extra benchmark files from the stratified sample")
    a = ap.parse_args()
    import records
    import test_identity_audio as g2
    import test_identity_cached as g1
    from anuvani_diar import Diarizer, _core
    failures = 0

    t0 = time.monotonic()
    d = Diarizer(keep_records=True, devices=a.devices.split(",") if a.devices else None, strict=True)
    print(f"1. self-test: every batch shape bit-identical on {len(d.nn.devices)} device(s) "
          f"({len(json.loads((PROJ / 'assets/nn_golden.json').read_text())['digests'])} shapes) — {time.monotonic() - t0:.1f} s",
          flush=True)

    t0 = time.monotonic()
    keys = [f"bench/{i['sid']}" for i in json.loads((PROJ / "artifacts/golden12_inputs.json").read_text())]
    keys += [k for k, i in g2.INPUTS.items() if i["set"] == "robust"]
    keys += ["long/long_1h"]
    keys += [f"{s}/{p.stem}" for s, p in records.sample(a.sample)]
    keys = [k for k in dict.fromkeys(keys) if k in g2.INPUTS]
    hours = sum(g2.INPUTS[k].get("seconds") or 0 for k in keys) / 3600
    bad = g2.compare(keys, d)
    failures += len(bad)
    print(f"2. engine vs raw-audio oracle: {len(keys) - len({b[0] for b in bad})}/{len(keys)} inputs identical "
          f"(records, RTTMs, errors; {hours:.2f} h of audio) — {time.monotonic() - t0:.1f} s", flush=True)
    for b in bad[:10]:
        print("   MISMATCH", b)

    t0 = time.monotonic()
    ok = 0
    sample = records.sample(40)
    for set_name, path in sample:
        rttm = g1.run(_core, set_name, path)
        want = (PROJ / "artifacts/oracle_cache/canonical" / set_name / "indic" / f"{path.stem}.rttm").read_bytes()
        ok += rttm.encode() == want
    failures += len(sample) - ok
    print(f"3. Rust post-processing vs cache oracle: {ok}/{len(sample)} RTTMs identical — {time.monotonic() - t0:.1f} s",
          flush=True)
    print("check-exact: PASS" if failures == 0 else f"check-exact: FAIL ({failures})")
    return 0 if failures == 0 else 1


if __name__ == "__main__":
    sys.exit(main())