"""Turn Nemotron 3 frame probabilities into speaker segments, replicating Omi's published decoding behaviour.

Omi's output differs from NVIDIA's plain 0.5 threshold by a slightly higher threshold and by bridging short
same-speaker pauses: always below `bridge_s`, and up to `bridge_clear_s` when no other speaker is active inside.
"""
from dataclasses import dataclass

import numpy as np

FRAME_S = 0.01


@dataclass(frozen=True)
class DecoderParams:
    threshold: float = 0.5
    bridge_s: float = 0.0
    bridge_clear_s: float = 0.0


def _runs(active: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    edges = np.diff(np.r_[False, active, False].astype(np.int8))
    return np.flatnonzero(edges == 1), np.flatnonzero(edges == -1)


def activity(probs: np.ndarray, params: DecoderParams) -> np.ndarray:
    """Boolean (frames, speakers) activity after thresholding and pause bridging."""
    raw = probs > params.threshold
    active = raw.copy()
    bridge, bridge_clear = round(params.bridge_s / FRAME_S), round(params.bridge_clear_s / FRAME_S)
    for speaker in range(raw.shape[1]):
        others = np.delete(raw, speaker, axis=1).any(axis=1)
        starts, ends = _runs(raw[:, speaker])
        for gap_start, gap_end in zip(ends[:-1], starts[1:]):
            length = gap_end - gap_start
            if length < bridge or (length < bridge_clear and not others[gap_start:gap_end].any()):
                active[gap_start:gap_end, speaker] = True
    return active


def decode(probs: np.ndarray, params: DecoderParams) -> list[dict]:
    """Segments in seconds; speakers that never activate are dropped, so the count stays automatic."""
    active = activity(probs, params)
    segments = []
    for speaker in range(active.shape[1]):
        starts, ends = _runs(active[:, speaker])
        segments.extend(
            {"start": round(float(s * FRAME_S), 2), "end": round(float(e * FRAME_S), 2), "speaker": f"spk_{speaker}"}
            for s, e in zip(starts, ends)
        )
    return segments