# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""The release's two networks, run with the release's model code, flags and batch shapes.

The networks are the pure-PyTorch copies in `seg/` (DiariZen's segmentation model) and `voiceprint.py` (pyannote's
WeSpeakerResNet34 with torchaudio's kaldi fbank), loaded from the bundle's safetensors (`tools/export_bundle.py`); no
DiariZen, pyannote, Lightning or torchaudio import. They run with the release's settings (FP32, TF32 off,
deterministic kernels, `CUBLAS_WORKSPACE_CONFIG=:4096:8`, 2 CPU threads), one copy per device, and their outputs are
bit-identical to the release's on every batch shape where the release's setup is reproduced (the start-up self-test
against assets/nn_golden.json). A *slot* is one device with one CUDA stream, driven by one engine worker thread: every
GPU gives the same bytes as the others of its kind, but two streams sharing one model on one GPU do not (about 1 batch
in 1,000 differed), so there is exactly one slot per device. What runs on a slot is always one of the release's own batch shapes:

  segment(x)                (B, 1, 256000) -> (B, 799, 16)   B = 8, the remainder (1-7), or 1 (padded last window)
  embed(crops, weights)     (B, 1, 256000), (B, 799) -> (B, 256)   the release's voiceprint call (B = 8 or 4)
  embed_repacked(...)       the voiceprint trunk once per window (batch 8), then pooling + the last layer at the
                            release's rows (2 windows x 4 speakers); bit-identical to `embed`
  resample(audio, sr)       `load_audio`'s torchaudio call (torchaudio's resample, copied in resample.py)

The engine (Rust) decides every batch; this module only executes them.
"""
import contextlib
import hashlib
import json
import os
import threading
import warnings
from pathlib import Path

import numpy as np

WIN, FRAMES, CLASSES, EMB = 256000, 799, 16, 256
FRAMES_SW = (-0.00753125, 0.025, 0.02)
MIN_NUM_SAMPLES = 400


# precision presets: (segmentation network, voiceprint trunk) compute types; "exact" is the release's eager FP32 path.
# "fp32c": FP32 matrix products with TF32 convolutions (cuDNN's TF32 switch on, cuBLAS's off).
# "ffn16": as fp32c, with the transformer's feed-forward blocks (two linear layers and a GELU) under FP16 autocast.
PRESETS = {"exact": ("exact", "exact"), "fp32": ("fp32", "fp32"), "mixed": ("fp32", "fp16"),
           "tf32conv": ("fp32c", "fp16"), "ffn16": ("ffn16", "fp16"), "tf32": ("tf32", "tf32"), "fp16": ("fp16", "fp16")}


# Byte-identical outputs need the exact software and hardware the release's outputs were produced with; on any other
# GPU or library version the exact lane computes the same FP32 networks, but the last bits of the floating-point
# results differ. Unless `strict`, it then runs if the self-test's outputs stay within these bounds of the release's
# own (assets/nn_reference.npz). The bounds were fixed before any measurement on other setups. They catch a broken
# installation; they are not the accuracy claim.
NEAR = dict(max_abs_dp=0.25, frame_agreement=0.99, min_cosine=0.995)


def resolve_devices(devices, torch):
    """The torch devices to run on, from `devices`:

    "auto" or None: every visible GPU, or the CPU when there is none;
    an int N: the first N visible GPUs;
    a list of GPU indices or device names, or a comma-separated string of them ([0, 2], "0,2", ["cuda:1"]);
    "cpu": the CPU only.
    """
    visible = torch.cuda.device_count()
    if devices is None or devices == "auto" or devices == "all":
        return [f"cuda:{i}" for i in range(visible)] or ["cpu"]
    if isinstance(devices, bool):
        raise ValueError("devices: 'auto', a number of GPUs, a list of GPUs, or 'cpu'")
    if isinstance(devices, int):
        if devices < 1:
            raise ValueError(f"devices={devices}: give a number of GPUs from 1, a list of GPUs, 'auto' or 'cpu'")
        if devices > visible:
            raise ValueError(f"devices={devices}: this machine has {visible} visible GPU(s)")
        return [f"cuda:{i}" for i in range(devices)]
    if isinstance(devices, str):
        devices = [d for d in devices.split(",") if d.strip()]
    out = []
    for d in devices:
        d = str(d).strip()
        if d.isdigit():
            d = f"cuda:{d}"
        if d.startswith("cuda") and int(d.partition(":")[2] or 0) >= visible:
            raise ValueError(f"{d}: this machine has {visible} visible GPU(s)")
        out.append(d)
    if not out:
        raise ValueError("devices: no device named")
    return out


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


def compare(name, a, b):
    """(report, within NEAR) for outputs `a` against reference outputs `b` of the self-test case `name`."""
    if name.startswith("segment"):
        b = b[:a.shape[0]]
        dp = float(np.abs(a - b).max())
        agree = float((a.argmax(-1) == b.argmax(-1)).mean())
        return dict(max_abs_dp=dp, frame_agreement=agree), dp <= NEAR["max_abs_dp"] and agree >= NEAR["frame_agreement"]
    na, nb = np.linalg.norm(a, axis=1), np.linalg.norm(b, axis=1)
    live = (na > 0) & (nb > 0)
    cos = float((np.sum(a * b, axis=1)[live] / (na[live] * nb[live])).min()) if live.any() else 1.0
    return dict(min_cosine=cos), cos >= NEAR["min_cosine"]


class TorchNN:
    def __init__(self, bundle, devices=None, workers_per_device=1, threads=2, golden=None, precision="exact",
                 strict=True):
        if precision not in PRESETS:
            raise ValueError(f"precision must be one of {sorted(PRESETS)}: 'exact' runs the release's FP32 networks, "
                             "the others faster reduced-precision variants")
        self.precision = precision
        self.strict = strict
        self.exactness = {}
        self._force_exact = False
        if workers_per_device != 1:
            raise ValueError("workers_per_device must be 1: two CUDA streams sharing one model on one GPU changed "
                             "about 1 batch in 1,000, so the engine runs one stream per GPU")
        os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
        if os.environ["CUBLAS_WORKSPACE_CONFIG"] != ":4096:8":
            raise RuntimeError("CUBLAS_WORKSPACE_CONFIG must be :4096:8 for byte-identical voiceprints")
        import copy
        import torch
        import torch.nn.functional as F
        from .seg import SegmentationNet
        from .voiceprint import VoiceprintNet
        self.torch, self.F = torch, F
        self.threads = threads
        torch.set_num_threads(threads)
        self.flag_resets = 0
        self._enforce_flags()
        bundle = Path(bundle)
        self.devices = [torch.device(d) for d in resolve_devices(devices, torch)]
        seg_cpu = SegmentationNet.from_pretrained(bundle / "segmentation")
        vp_cpu = VoiceprintNet.from_pretrained(bundle / "voiceprint")
        self.nets = []
        for dev in self.devices:
            # each device gets its own copy, moved from host memory (an exact byte copy) and checked
            seg, vp = copy.deepcopy(seg_cpu).to(dev).eval(), copy.deepcopy(vp_cpu).to(dev).eval()
            for net in (seg, vp):
                if any(bool(torch.isnan(t).any()) for t in list(net.parameters()) + list(net.buffers()) if t.is_floating_point()):
                    raise RuntimeError(f"NaN weights after moving the networks to {dev}")
            self.nets.append((seg, vp))
        del seg_cpu, vp_cpu
        self._fast = [self._compile_fast(seg, vp) for seg, vp in self.nets] if precision != "exact" else None
        self.slots = []
        for di, dev in enumerate(self.devices):
            for _ in range(workers_per_device if dev.type == "cuda" else 1):
                stream = torch.cuda.Stream(device=dev) if dev.type == "cuda" else None
                self.slots.append((di, dev, stream))
        self._golden = json.loads(Path(golden).read_text()) if golden else None
        ref = Path(golden).with_name("nn_reference.npz") if golden else None
        self._reference = dict(np.load(ref)) if ref is not None and ref.is_file() else None
        self._lock = threading.Lock()
        self.checked = []
        # warm every shape on every slot before any concurrent use
        for s in range(len(self.slots)):
            self.self_test(s)

    def _enforce_flags(self):
        """The release's numeric flags (load_bundle). They are process-global, so any later import or user code may
        change them; every network call re-asserts them (and counts how often it had to)."""
        t = self.torch
        exact = self.precision == "exact" or self._force_exact
        types = PRESETS[self.precision]
        tf32 = not exact and "tf32" in types
        conv_tf32 = tf32 or (not exact and ("fp32c" in types or "ffn16" in types))
        want = (exact, tf32, conv_tf32, not exact, exact)
        have = (t.are_deterministic_algorithms_enabled(), t.backends.cuda.matmul.allow_tf32, t.backends.cudnn.allow_tf32,
                t.backends.cudnn.benchmark, t.backends.cudnn.deterministic)
        if have == want:
            return
        self.flag_resets += 1
        t.use_deterministic_algorithms(exact)
        t.backends.cuda.matmul.allow_tf32 = tf32
        t.backends.cudnn.allow_tf32 = conv_tf32
        t.backends.cudnn.benchmark = not exact
        t.backends.cudnn.deterministic = exact

    # ------------------------------------------------------------------ info
    def info(self):
        t = self.torch
        return dict(torch=t.__version__, cuda=t.version.cuda, cudnn=t.backends.cudnn.version(),
                    devices=[str(d) for d in self.devices],
                    gpus=[t.cuda.get_device_name(d) if d.type == "cuda" else "cpu" for d in self.devices],
                    slots=[(di, str(dev)) for di, dev, _ in self.slots], frames_sw=list(FRAMES_SW),
                    min_num_samples=MIN_NUM_SAMPLES, threads=self.threads, flag_resets=self.flag_resets,
                    precision=self.precision, exactness={str(k): v for k, v in self.exactness.items()})

    def slot_devices(self):
        return [di for di, _, _ in self.slots]

    @contextlib.contextmanager
    def _on(self, slot):
        self._enforce_flags()
        di, dev, stream = self.slots[slot]
        t = self.torch
        with contextlib.ExitStack() as st:
            if stream is not None:
                st.enter_context(t.cuda.device(dev))
                st.enter_context(t.cuda.stream(stream))
            st.enter_context(t.inference_mode())
            yield self.nets[di], dev

    # ------------------------------------------------------------------ the release's calls
    def segment(self, slot, x):
        """`Inference.infer` with the softmax conversion of `load_bundle`. x: (B, 1, 256000) float32 numpy."""
        return self.segment_t(slot, lambda dev: self.torch.from_numpy(x).to(dev))

    def embed(self, slot, crops, weights):
        """`PyannoteAudioPretrainedSpeakerEmbedding.__call__(waveforms, masks)`."""
        return self.embed_t(slot, lambda dev: self.torch.from_numpy(crops).to(dev), weights)

    def embed_repacked(self, slot, windows, idx, weights):
        """Voiceprints of H reference batches from one trunk pass over 8 distinct windows.

        windows (8, 1, 256000) float32 (zero rows pad a short block); idx (H, 8) int64, the trunk row of every
        reference row; weights (H, 8, 799) float32. ResNet.forward split at the pooling layer."""
        return self.embed_repacked_t(slot, lambda dev: self.torch.from_numpy(windows).to(dev), idx, weights)

    # The same three calls with the windows built on the device by `make(dev)` (a tensor (B, 1, 256000) float32,
    # contiguous): from numpy above, or from a span of samples (`windows_from_span`) in the Triton model.
    def windows_from_span(self, span, offsets, valid):
        """`make(dev)` for windows given as one span: row j = span[offsets[j]:offsets[j] + valid[j]], then zeros."""
        t = self.torch

        def make(dev):
            s = t.from_numpy(span).to(dev)
            x = t.zeros((len(offsets), 1, WIN), dtype=t.float32, device=dev)
            for j, (o, v) in enumerate(zip(offsets.tolist(), valid.tolist())):
                if v:
                    x[j, 0, :v] = s[o:o + v]
            return x
        return make

    # ------------------------------------------------------------------ the fast lane (not byte-identical)
    FAST_BATCH = 8

    def _compile_fast(self, seg, vp):
        """The two networks under torch.compile at one shape (batches padded to 8 windows), each in the compute type
        of the preset (PRESETS): fp32 (TF32 off), tf32, or fp16 autocast. fbank, pooling, the last layer and the softmax
        stay in FP32 (fbank power reaches ~1e14, above FP16's range)."""
        t = self.torch
        r = vp.resnet
        seg_type, trunk_type = PRESETS[self.precision]
        if seg_type == "ffn16":
            # a copy, so that the exact path (the self-test, strict mode) keeps the unchanged modules
            import copy
            from .seg.wav2vec2.components import FeedForward
            seg = copy.deepcopy(seg)
            for m in seg.modules():
                if isinstance(m, FeedForward):
                    def fwd(x, f=m.forward):
                        with t.autocast("cuda", dtype=t.float16):
                            y = f(x)
                        return y.float()
                    m.forward = fwd

        def seg_c(x):
            with t.autocast("cuda", dtype=t.float16, enabled=seg_type == "fp16"):
                out = seg(x)
            return t.softmax(out.float(), -1)

        def trunk_c(fb):
            with t.autocast("cuda", dtype=t.float16, enabled=trunk_type == "fp16"):
                h = fb.permute(0, 2, 1).unsqueeze(1)
                h = self.F.relu(r.bn1(r.conv1(h)))
                h = r.layer4(r.layer3(r.layer2(r.layer1(h))))
            return h.float()
        return t.compile(seg_c, dynamic=False), t.compile(trunk_c, dynamic=False)

    def _pad(self, x):
        b = x.shape[0]
        if b == self.FAST_BATCH:
            return x, b
        pad = self.torch.zeros((self.FAST_BATCH - b,) + tuple(x.shape[1:]), dtype=x.dtype, device=x.device)
        return self.torch.cat([x, pad]), b

    def _fast_on(self, slot):
        return self._fast is not None and not self._force_exact

    def _trunk(self, slot, emb, x):
        """The voiceprint trunk's output for windows x (B, 1, 256000): exact (the release's ResNet.forward up to the
        pooling layer) or the fast lane's compiled FP16 trunk."""
        t, F = self.torch, self.F
        r = emb.resnet
        fb = emb.compute_fbank(x)
        if self._fast_on(slot):
            out = []
            for i in range(0, fb.shape[0], self.FAST_BATCH):
                part, b = self._pad(fb[i:i + self.FAST_BATCH])
                out.append(self._fast[self.slots[slot][0]][1](part)[:b])
            return t.cat(out)
        x = fb.permute(0, 2, 1)
        x = x.unsqueeze_(1)
        out = F.relu(r.bn1(r.conv1(x)))
        out = r.layer1(out)
        out = r.layer2(out)
        out = r.layer3(out)
        return r.layer4(out)

    def segment_t(self, slot, make):
        with self._on(slot) as ((seg, _), dev):
            x = make(dev)
            if self._fast_on(slot):
                out = []
                for i in range(0, x.shape[0], self.FAST_BATCH):
                    part, b = self._pad(x[i:i + self.FAST_BATCH])
                    out.append(self._fast[self.slots[slot][0]][0](part)[:b])
                return self.torch.cat(out).cpu().numpy()
            out = seg(x)
            return self.torch.softmax(out, -1).cpu().numpy()

    def embed_t(self, slot, make, weights):
        t = self.torch
        with self._on(slot) as ((_, emb), dev), warnings.catch_warnings():
            warnings.simplefilter("ignore")
            if self._fast_on(slot):
                r = emb.resnet
                trunk = self._trunk(slot, emb, make(dev))
                return r.seg_1(r.pool(trunk, weights=t.from_numpy(weights).to(dev))).cpu().numpy()
            return emb(make(dev), weights=t.from_numpy(weights).to(dev)).cpu().numpy()

    def embed_repacked_t(self, slot, make, idx, weights):
        t, F = self.torch, self.F
        with self._on(slot) as ((_, emb), dev), warnings.catch_warnings():
            warnings.simplefilter("ignore")
            r = emb.resnet
            trunk = self._trunk(slot, emb, make(dev))
            res = []
            for h in range(idx.shape[0]):
                rows = trunk[t.from_numpy(idx[h]).to(dev)]
                res.append(r.seg_1(r.pool(rows, weights=t.from_numpy(weights[h]).to(dev))).cpu().numpy())
            return np.concatenate(res) if res else np.zeros((0, EMB), np.float32)

    def resample(self, audio, sr):
        """`load_audio`'s resampling, with the calling thread's CPU thread count set as the release sets it."""
        from .resample import resample
        t = self.torch
        t.set_num_threads(self.threads)
        return resample(t.from_numpy(audio).unsqueeze(0), sr, 16000).squeeze(0).numpy().astype(np.float32)

    # ------------------------------------------------------------------ guard
    def self_test(self, slot):
        """Run every shape on this slot and compare with the release's own outputs (assets/nn_golden.json): equal
        bytes, or, where those cannot be reproduced and unless `strict`, within NEAR of assets/nn_reference.npz."""
        g = self._golden
        res = {}
        quick = self.slots[slot][1].type == "cpu"
        self._force_exact = True
        try:
            exact = {name: fn() for name, fn in golden_cases(self, slot, quick=quick)}
        finally:
            self._force_exact = False
        for name, out in exact.items():
            res[name] = _digest(out)
        if g is not None:
            bad = [k for k in g["digests"] if res.get(k) != g["digests"][k]]
            if not bad:
                self.exactness[slot] = dict(status="byte-identical")
            elif self.strict or self._reference is None:
                raise RuntimeError(f"slot {slot} ({self.slots[slot][1]}): outputs differ from the release for {bad}"
                                   + ("" if self.strict else " (no assets/nn_reference.npz to check them against)"))
            else:
                near, ok = {}, True
                for name, out in exact.items():
                    near[name], good = compare(name, out, self._reference[name if not name.startswith("segment") else "segment_b8"])
                    ok = ok and good
                if not ok:
                    raise RuntimeError(f"slot {slot} ({self.slots[slot][1]}): outputs are outside the near-identical bounds "
                                       f"{NEAR} of the release's: {near}")
                worst = dict(max_abs_dp=max(v["max_abs_dp"] for k, v in near.items() if k.startswith("segment")),
                             frame_agreement=min(v["frame_agreement"] for k, v in near.items() if k.startswith("segment")),
                             min_cosine=min(v["min_cosine"] for k, v in near.items() if not k.startswith("segment")))
                self.exactness[slot] = dict(status="near-identical", **worst)
                warnings.warn(f"anuvani_diar: on {self.slots[slot][1]} the networks run in FP32 and agree with the "
                              f"release's outputs within the near-identical bounds, but not to the last bit (that needs "
                              f"the setup the release's outputs were produced with); self-test: {worst}", stacklevel=3)
        if self._fast is not None:
            # the fast lanes' start-up bounds: posteriors within 0.25, frame decisions >= 99% equal,
            # voiceprint cosine >= 0.995 on these synthetic inputs (noise and a zero-padded, mostly silent window)
            fast = {name: fn() for name, fn in golden_cases(self, slot)}
            dev_report = {}
            for name in exact:
                a, b = exact[name], fast[name]
                if name.startswith("segment"):
                    dp = float(np.abs(a - b).max())
                    agree = float((a.argmax(-1) == b.argmax(-1)).mean())
                    ok = dp <= 0.25 and agree >= 0.99
                    dev_report[name] = dict(max_abs_dp=dp, frame_agreement=agree)
                else:
                    na, nb = np.linalg.norm(a, axis=1), np.linalg.norm(b, axis=1)
                    live = (na > 0) & (nb > 0)
                    cos = float((np.sum(a * b, axis=1)[live] / (na[live] * nb[live])).min()) if live.any() else 1.0
                    ok = cos >= 0.995
                    dev_report[name] = dict(min_cosine=cos)
                if not ok:
                    raise RuntimeError(f"slot {slot} ({self.slots[slot][1]}): the fast lane is outside its start-up bounds for "
                                       f"{name}: {dev_report[name]}")
            self.fast_report = getattr(self, "fast_report", {})
            self.fast_report[slot] = dev_report
        with self._lock:
            self.checked.append(slot)
        return res


def golden_inputs():
    """Deterministic inputs of the self-test: seeded noise windows and seeded binary masks."""
    rng = np.random.default_rng(20261007)
    win = (rng.standard_normal((8, 1, WIN)) * 0.1).astype(np.float32)
    # a zero-padded window
    win[3, 0, 100000:] = 0
    masks = (rng.random((8, 4, FRAMES)) < 0.4).astype(np.float32)
    masks[1, 2] = 0
    return win, masks


def golden_cases(nn, slot, quick=False):
    """(name, thunk) for every batch shape the engine sends to a slot (quick: one window and one voiceprint batch, for
    a CPU, where the full set takes minutes)."""
    win, masks = golden_inputs()
    if quick:
        return [("segment_b1", lambda: nn.segment(slot, win[:1].copy())),
                ("embed_b4", lambda: nn.embed(slot, np.repeat(win[2:3], 4, 0), masks[2].copy()))]
    cases = [(f"segment_b{b}", (lambda b=b: nn.segment(slot, win[:b].copy()))) for b in range(1, 9)]
    pair = np.concatenate([np.repeat(win[0:1], 4, 0), np.repeat(win[1:2], 4, 0)])
    w8 = np.concatenate([masks[0], masks[1]])
    cases.append(("embed_b8", lambda: nn.embed(slot, pair, w8)))
    cases.append(("embed_b4", lambda: nn.embed(slot, np.repeat(win[2:3], 4, 0), masks[2].copy())))
    idx = np.array([[0, 0, 0, 0, 1, 1, 1, 1], [2, 2, 2, 2, 3, 3, 3, 3]], np.int64)
    wr = np.stack([w8, np.concatenate([masks[2], masks[3]])])
    cases.append(("embed_repacked_h2", lambda: nn.embed_repacked(slot, win.copy(), idx, wr)))
    return cases