# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""GPU probes (batching, voiceprints, devices): may the network side change anything without changing a single output
bit?

All models are built by the frozen release's own `load_bundle` (canonical flags: FP32, TF32 off, deterministic
kernels, CUBLAS_WORKSPACE_CONFIG=:4096:8), on real IndicDiarBench audio and real cached masks. Every comparison is
bitwise on the float32 outputs. Decision rules (written before running):
  batching     a row's segmentation output must not change when other rows' content or its position in the batch
               change (fixed batch shape); batch-size dependence is reported separately. Cross-file batching is
               allowed only if every row is bit-identical across all tested batch sizes.
  voiceprints  the repacked voiceprint trunk (fbank + ResNet trunk once per window, in batches of 8 distinct windows;
               pooling and seg_1 at the reference rows and shapes) is allowed only if every final embedding is
               bit-identical.
  devices      concurrent CUDA streams in one process, and each GPU, must reproduce GPU 0's sequential bits; with
               --occupy the reference is repeated with that fraction (for example 0.95) of GPU memory occupied before
               the first forward (cuDNN engine fallback).

  source env/canonical.sh && CUDA_VISIBLE_DEVICES=0 python -s tools/probes.py batching --out artifacts/probes
"""
import argparse
import hashlib
import json
import os
import sys
import threading
import time
from pathlib import Path

import numpy as np

PROJ = Path(__file__).resolve().parents[1]
REL = PROJ.parent / "anuvani_diar_v1"
SID = PROJ.parents[1]
sys.path.insert(0, str(PROJ / "tests"))
WIN, STEP = 256000, 25600
# golden files (3 languages)
FILES = ["bengali_009", "hindi_033", "kannada_035"]


def load_ref():
    import refpkg
    ref = refpkg.load()
    from anuvani_diar_ref.frontend import load_bundle
    return ref, load_bundle


def audio(sid):
    import soundfile as sf
    p = next((SID / "benchmark_data/indicdiar_full").glob(f"*/{sid}.flac"))
    a, sr = sf.read(str(p), dtype="float32")
    assert sr == 16000 and a.ndim == 1
    return a


def windows(a, n):
    """First n reference windows (chunk i starts at i * 25600; zero-padded if short)."""
    out = np.zeros((n, 1, WIN), np.float32)
    for i in range(n):
        seg = a[i * STEP:i * STEP + WIN]
        out[i, 0, :len(seg)] = seg
    return out


def masks_for(sid):
    import records
    p = next(p for s, p in records.all_records() if p.stem == sid)
    return np.load(p)["masks"]


def used_weights(masks):
    """pyannote get_embeddings(exclude_overlap=True) mask choice: clean frames if more than 2, else the full mask."""
    m = masks.astype(np.float64)
    clean = m * (m.sum(axis=2, keepdims=True) < 2)
    w = np.where((clean.sum(axis=1, keepdims=True) > 2), clean, m)
    # (C, 799, 4)
    return np.nan_to_num(w).astype(np.float32)


def sha(t):
    return hashlib.sha256(t.detach().cpu().contiguous().numpy().tobytes()).hexdigest()[:16]


def seg_forward(model, x, torch):
    with torch.inference_mode():
        return torch.softmax(model(torch.from_numpy(x).to("cuda")), -1).cpu()


def emb_parts(emb_model, torch, F):
    r = emb_model.resnet

    def trunk(w):
        fb = emb_model.compute_fbank(w)
        x = fb.permute(0, 2, 1)
        x = x.unsqueeze_(1)
        out = F.relu(r.bn1(r.conv1(x)))
        return r.layer4(r.layer3(r.layer2(r.layer1(out))))

    def head(out, weights):
        return r.seg_1(r.pool(out, weights=weights))
    return trunk, head


def batching(args):
    import torch
    import torch.nn.functional as F
    ref, load_bundle = load_ref()
    model = load_bundle(REL / "bundle", "cuda:0", 8, 2)
    seg = model["pipeline"]._segmentation.model
    emb = model["pipeline"]._embedding.model_
    trunk, head = emb_parts(emb, torch, F)
    res = dict(torch=torch.__version__, cudnn=torch.backends.cudnn.version(), gpu=torch.cuda.get_device_name(0),
               deterministic=torch.are_deterministic_algorithms_enabled(), segmentation={}, voiceprints={})
    rng = np.random.default_rng(0)
    for sid in FILES:
        a = audio(sid)
        W = windows(a, 24)
        # reference: a full batch of 8, in order
        base = seg_forward(seg, W[:8], torch)
        seg_res = res["segmentation"][sid] = {}
        perm = rng.permutation(8)
        seg_res["position_change"] = bool(torch.equal(seg_forward(seg, W[:8][perm], torch)[np.argsort(perm)], base))
        # other rows replaced by other windows
        mixed = W[:8].copy(); mixed[1::2] = W[8:16][1::2]
        seg_res["other_rows_content_change"] = bool(torch.equal(seg_forward(seg, mixed, torch)[0::2], base[0::2]))
        zero = W[:8].copy(); zero[1::2] = 0.0
        seg_res["other_rows_zeroed"] = bool(torch.equal(seg_forward(seg, zero, torch)[0::2], base[0::2]))
        seg_res["batch_size_identical"] = {}
        for b in (1, 2, 3, 4, 5, 6, 7, 9, 12, 16, 24):
            Wb = W[:b] if b <= 24 else W
            ob = seg_forward(seg, Wb, torch)
            k = min(b, 8)
            seg_res["batch_size_identical"][b] = bool(torch.equal(ob[:k], base[:k]))
        # voiceprints, reference rows = 2 chunks x 4 speakers (identical crops), real weights
        m = used_weights(masks_for(sid))
        vp_res = res["voiceprints"][sid] = {}
        with torch.inference_mode():
            # three reference batches
            for start in (0, 2, 6):
                crops = torch.from_numpy(np.repeat(W[start:start + 2], 4, axis=0)).to("cuda")
                wts = torch.from_numpy(np.stack([m[start + i // 4, :, i % 4] for i in range(8)])).to("cuda")
                # what the release computes
                ref_emb = emb(crops, weights=wts)
                # 8 distinct windows, same shape
                distinct = torch.from_numpy(W[start:start + 8]).to("cuda")
                T = trunk(distinct)
                # gather the two chunks' rows, x4 each
                Tg = T[[0, 0, 0, 0, 1, 1, 1, 1]].contiguous()
                rep = head(Tg, wts)
                T_ref = trunk(crops)
                vp_res[f"batch_{start}"] = dict(embedding_identical=bool(torch.equal(rep, ref_emb)),
                                              trunk_rows_identical=bool(torch.equal(T_ref[0], T[0]) and torch.equal(T_ref[4], T[1])),
                                              max_abs_diff=float((rep - ref_emb).abs().max()))
            # slot positions: the same window at each of the 8 slots of a distinct batch
            pos_ok = True
            for slot in range(8):
                batch = W[8:16].copy(); batch[slot] = W[0]
                T = trunk(torch.from_numpy(batch).to("cuda"))
                T0 = trunk(torch.from_numpy(np.repeat(W[0:1], 8, axis=0)).to("cuda"))
                pos_ok &= bool(torch.equal(T[slot], T0[0]))
            vp_res["trunk_all_slot_positions"] = pos_ok
            # tail: the zero-padded last chunk and digital silence
            # a short, zero-padded window
            tail = windows(a[: 5 * STEP + 70000], 6)[-1:]
            silent = np.zeros((1, 1, WIN), np.float32)
            for name, w in (("padded_tail", tail), ("digital_silence", silent)):
                ref8 = trunk(torch.from_numpy(np.repeat(w, 8, axis=0)).to("cuda"))
                mix = W[8:16].copy(); mix[3] = w[0]
                got = trunk(torch.from_numpy(mix).to("cuda"))
                vp_res[f"trunk_{name}"] = bool(torch.equal(got[3], ref8[0]))
    out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
    (out / "batching.json").write_text(json.dumps(res, indent=1))
    print(json.dumps(res, indent=1))


def devices_reference(args):
    """Sequential reference digests on the visible GPU (also with most of its memory occupied, when --occupy is set)."""
    import torch
    import torch.nn.functional as F
    if args.occupy:
        free, total = torch.cuda.mem_get_info()
        # held for the run
        hog = torch.empty(int(free * args.occupy) // 4, dtype=torch.float32, device="cuda")
    ref, load_bundle = load_ref()
    model = load_bundle(REL / "bundle", "cuda:0", 8, 2)
    seg = model["pipeline"]._segmentation.model
    emb = model["pipeline"]._embedding.model_
    dig = {}
    for sid in FILES:
        W = windows(audio(sid), 16)
        m = used_weights(masks_for(sid))
        with torch.inference_mode():
            dig[f"{sid}/seg"] = sha(seg_forward(seg, W[:8], torch))
            crops = torch.from_numpy(np.repeat(W[0:2], 4, axis=0)).to("cuda")
            wts = torch.from_numpy(np.stack([m[i // 4, :, i % 4] for i in range(8)])).to("cuda")
            dig[f"{sid}/emb"] = sha(emb(crops, weights=wts))
    if args.streams:
        # two threads, each on its own CUDA stream, running the same batches concurrently
        outs = {}

        def work(name, sid):
            W = windows(audio(sid), 16)
            s = torch.cuda.Stream()
            with torch.cuda.stream(s), torch.inference_mode():
                for _ in range(3):
                    o = torch.softmax(seg(torch.from_numpy(W[:8]).to("cuda", non_blocking=False)), -1)
                s.synchronize()
                outs[name] = sha(o)
        ts = [threading.Thread(target=work, args=(f"t{i}", FILES[i % len(FILES)])) for i in range(4)]
        [t.start() for t in ts]; [t.join() for t in ts]
        dig["streams"] = {k: v == dig[f"{FILES[int(k[1:]) % len(FILES)]}/seg"] for k, v in outs.items()}
    gpu = os.environ.get("CUDA_VISIBLE_DEVICES", "?")
    out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
    tag = f"devices_gpu{gpu}" + (f"_occupy{args.occupy}" if args.occupy else "") + ("_streams" if args.streams else "")
    (out / f"{tag}.json").write_text(json.dumps(dig, indent=1))
    print(tag, json.dumps(dig))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("what", choices=["batching", "devices"])
    ap.add_argument("--out", default=str(PROJ / "artifacts/probes"))
    ap.add_argument("--occupy", type=float, default=0.0)
    ap.add_argument("--streams", action="store_true")
    a = ap.parse_args()
    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")
    batching(a) if a.what == "batching" else devices_reference(a)


if __name__ == "__main__":
    main()