# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Deterministic stand-ins for the two networks, for CPU tests of everything around them.

Each output row is a pure function of that row's input (as measured for the real networks at the release's batch
shapes), so the release's `analyse` and the engine must build byte-identical records from them whatever the batching.
`FakeNN` is the engine-side backend; `patch_release` installs the same functions in the release's pipeline.
"""
import numpy as np

WIN, FRAMES, CLASSES, EMB = 256000, 799, 16, 256
FRAMES_SW = [-0.00753125, 0.025, 0.02]
_F = np.arange(FRAMES, dtype=np.float64)[:, None]
_K = np.arange(CLASSES)[None, :]
_D = np.arange(EMB, dtype=np.float64)


def fake_segment(x):
    x = np.asarray(x)
    out = np.empty((x.shape[0], FRAMES, CLASSES), np.float32)
    for r in range(x.shape[0]):
        v = x[r, 0, :: WIN // 32].astype(np.float64)
        a = float(np.abs(x[r, 0, ::997]).astype(np.float64).sum())
        logits = np.sin(_F * (0.011 + (a % 1.0) * 0.004) + _K * (1.3 + v[_K % 32] * 3.0)) * 5.0
        # silence -> "no speaker"
        logits[:, 0] += 6.0 * (a < 1e-3)
        e = np.exp(logits - logits.max(1, keepdims=True))
        out[r] = (e / e.sum(1, keepdims=True)).astype(np.float32)
    return out


def fake_embed(crops, weights):
    crops, weights = np.asarray(crops), np.asarray(weights)
    out = np.empty((crops.shape[0], EMB), np.float32)
    for r in range(crops.shape[0]):
        c = crops[r, 0, ::1000].astype(np.float64)
        w = weights[r].astype(np.float64)
        s = w.sum()
        out[r] = (np.cos(c * 7.0 + _D * 0.1) + np.sin(_D * (1.0 + (s % 5.0))) + w[::4][:EMB].sum() * 1e-3).astype(np.float32)
    return out


class FakeNN:
    """Engine backend with the stand-in networks; records every call's shape."""

    def __init__(self, devices=(0, 0, 1)):
        self._slots = list(devices)
        self.calls = []

    def slot_devices(self):
        return self._slots

    def info(self):
        return dict(frames_sw=FRAMES_SW)

    def segment(self, slot, x):
        self.calls.append(("segment", x.shape))
        return fake_segment(x)

    def embed(self, slot, crops, weights):
        self.calls.append(("embed", crops.shape, weights.shape))
        return fake_embed(crops, weights)

    def embed_repacked(self, slot, windows, idx, weights):
        self.calls.append(("embed_repacked", windows.shape, idx.shape, weights.shape))
        return np.concatenate([fake_embed(windows[idx[h]], weights[h]) for h in range(idx.shape[0])])

    def resample(self, audio, sr):
        import torch
        import torchaudio.functional as F
        torch.set_num_threads(2)
        return F.resample(torch.from_numpy(audio).unsqueeze(0), sr, 16000).squeeze(0).numpy().astype(np.float32)


def patch_release(model):
    """Replace the networks inside a pipeline built by the release's `load_bundle`."""
    pipe = model["pipeline"]
    pipe._segmentation.infer = lambda chunks, soft=False: fake_segment(chunks.numpy())

    class Emb:
        sample_rate = 16000
        min_num_samples = 400
        dimension = EMB
        metric = "cosine"

        def __call__(self, waveforms, masks=None):
            return fake_embed(waveforms.numpy(), masks.numpy())

    pipe._embedding = Emb()