# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Trace pyannote's own front-end batching (`Inference.slide`, `SpeakerDiarization.get_embeddings`) with both
networks replaced by recorders, on CPU, and print the batches as JSON. Used by tests/test_front.py.

The waveform carries its own sample index (sample k has value k + 1, exact in float32 below 2^24), so each
recorded window row tells which samples it holds; zero padding stays 0.

  python -s tests/ref_front_trace.py LENGTHS_JSON MASK_SEED OUT_JSON
"""
import json
import os
import sys
from pathlib import Path

# never touch a GPU from the tests
os.environ["CUDA_VISIBLE_DEVICES"] = ""

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
import refpkg                                               # noqa: E402

refpkg.load()
import torch                                                # noqa: E402
from pyannote.core import SlidingWindowFeature              # noqa: E402
from anuvani_diar_ref.frontend import load_bundle             # noqa: E402

REL = Path(__file__).resolve().parents[2] / "anuvani_diar_v1"


def describe(rows):
    """For each (1, WIN) row: first sample index, number of non-padding samples, padding all zero."""
    out = []
    for r in rows:
        r = r.reshape(-1)
        nz = np.flatnonzero(r)
        valid = int(nz[-1]) + 1 if len(nz) else 0
        first = int(r[0]) - 1 if valid else -1
        ok = bool(valid == 0 or np.array_equal(r[:valid], np.arange(first + 1, first + 1 + valid, dtype=np.float32)))
        out.append([first, valid, ok and not r[valid:].any()])
    return out


def main():
    lengths = json.loads(sys.argv[1])
    rng = np.random.default_rng(int(sys.argv[2]))
    model = load_bundle(REL / "bundle", "cpu", 8, 2)
    pipe = model["pipeline"]
    # the real voiceprint model, on CPU
    real_min = pipe._embedding.min_num_samples
    seg_rec, emb_rec = [], []

    def fake_infer(chunks, soft=False):
        x = chunks.numpy()
        seg_rec.append(dict(B=int(x.shape[0]), rows=describe(x)))
        return np.zeros((x.shape[0], 799, 16), np.float32)

    class FakeEmbedding:
        sample_rate = 16000
        min_num_samples = real_min
        dimension = 256
        metric = "cosine"

        def __call__(self, waveforms, masks=None):
            emb_rec.append(dict(B=int(waveforms.shape[0]), rows=describe(waveforms.numpy()),
                                masks=masks.numpy().astype(np.uint8).tolist()))
            return np.zeros((waveforms.shape[0], 256), np.float32)

    pipe._segmentation.infer = fake_infer
    pipe._embedding = FakeEmbedding()
    result = dict(min_num_samples=int(real_min), receptive_field=[model["frames"].start, model["frames"].duration,
                                                                  model["frames"].step],
                  mapping=model["mapping"].tolist(), runs=[])
    for n in lengths:
        assert n < 2 ** 24
        audio = np.arange(1, n + 1, dtype=np.float32)
        file = {"waveform": torch.from_numpy(audio).unsqueeze(0), "sample_rate": 16000}
        seg_rec.clear(); emb_rec.clear()
        seg = pipe.get_segmentations(file, soft=False)
        c = seg.data.shape[0]
        density = rng.choice([0.05, 0.3, 0.6, 0.95])
        masks = (rng.random((c, 799, 4)) < density).astype(np.float32)
        # exact clean-frame counts at the threshold
        if c > 1:
            masks[1] = 0
            # 2 clean frames: falls back to the full mask
            masks[1, :2, 0] = 1
            # 3 clean frames: uses the clean mask
            masks[1, 10:13, 1] = 1
            # overlapped frames
            masks[1, 13:20, 1:3] = 1
        sw = seg.sliding_window
        emb = pipe.get_embeddings(file, SlidingWindowFeature(masks, sw), exclude_overlap=True)
        result["runs"].append(dict(n=n, chunks=c, sw=[sw.start, sw.duration, sw.step], seg=list(seg_rec),
                                   emb=list(emb_rec), masks=masks.astype(np.uint8).tolist(),
                                   emb_shape=list(emb.shape)))
    Path(sys.argv[3]).write_text(json.dumps(result))


if __name__ == "__main__":
    main()