"""Postprocessing for DiariZen segmentation output."""

import numpy as np


def log_probs_to_probs(log_probs: np.ndarray) -> np.ndarray:
    """Convert log probabilities to probabilities via softmax.

    Args:
        log_probs: (1, frames, 11) log-softmax output.

    Returns:
        (1, frames, 11) probability distribution.
    """
    max_val = log_probs.max(axis=-1, keepdims=True)
    exp_vals = np.exp(log_probs - max_val)
    return exp_vals / exp_vals.sum(axis=-1, keepdims=True)


def top_speakers_at_frame(
    log_probs: np.ndarray,
    frame_idx: int,
    top_k: int = 3,
) -> list[tuple[int, float]]:
    """Get top-k speaker class indices and their log-probabilities at a frame.

    Args:
        log_probs: (1, frames, 11) output.
        frame_idx: Frame index.
        top_k: Number of top classes.

    Returns:
        List of (class_index, log_probability) tuples.
    """
    frame = log_probs[0, frame_idx]
    top_indices = np.argsort(frame)[-top_k:][::-1]
    return [(int(i), float(frame[i])) for i in top_indices]