# SPDX-License-Identifier: AGPL-3.0-only
# Copyright (c) 2026 sol pbc

"""Overlap-fraction inference for speaker-attribution flywheel gating."""

from __future__ import annotations

import logging
import time
from pathlib import Path
from typing import TYPE_CHECKING, NamedTuple, Sequence

from solstone.apps.speakers.encoder_config import (
    DIARIZE_MIN_OVERLAP,
    SLOT_ACTIVE_MIN_SHARE,
    SPEAKER_EVIDENCE_MULTI_MIN,
    SPEAKER_EVIDENCE_SINGLE_MAX,
)
from solstone.observe.utils import SAMPLE_RATE

if TYPE_CHECKING:
    import numpy as np
    import onnxruntime as ort

logger = logging.getLogger(__name__)

WINDOW_S = 10
STRIDE_S = 5
FRAMES_PER_WINDOW = 589
OVERLAP_CLASSES = (4, 5, 6)

# Tighter stride used when sharing pyannote output with the diarizer
_DIARIZE_STRIDE_S = 2

_overlap_session: ort.InferenceSession | None = None


class SpeakerWindowStats(NamedTuple):
    speech_frames: int
    active_slot_count: int
    overlap_frames: int


class SpeakerEvidenceDecision(NamedTuple):
    speaker_evidence: str
    multi_window_fraction: float
    mean_window_overlap_share: float


class OverlapInferenceResult(NamedTuple):
    overlap_fraction: float
    avg_log_probs: np.ndarray
    window_stats: tuple[SpeakerWindowStats, ...]


def _speaker_window_stats(log_probs: np.ndarray) -> SpeakerWindowStats:
    """Derive raw speaker-evidence counts from one truncated pyannote window."""
    import numpy as np

    argmax = log_probs.argmax(axis=-1)
    counts = np.bincount(argmax, minlength=7)
    speech_frames = int(counts[1:].sum())
    if speech_frames == 0:
        return SpeakerWindowStats(
            speech_frames=0,
            active_slot_count=0,
            overlap_frames=0,
        )

    active_slot_count = int(
        (counts[1:4] / speech_frames >= SLOT_ACTIVE_MIN_SHARE).sum()
    )
    return SpeakerWindowStats(
        speech_frames=speech_frames,
        active_slot_count=active_slot_count,
        overlap_frames=int(counts[list(OVERLAP_CLASSES)].sum()),
    )


def decide_speaker_evidence(
    overlap_fraction: float, window_stats: Sequence[SpeakerWindowStats]
) -> SpeakerEvidenceDecision:
    """Decide whether local diarization should engage for this segment.

    ``window overlap share`` is overlap-class frames divided by speech frames
    within a single window, from that window's raw argmax. It is explicitly not
    the segment-level ``overlap_fraction``, which argmaxes averaged log-probs
    across overlapping windows.

    ``mean_window_overlap_share`` is averaged over speech-bearing windows only.
    ``DIARIZE_MIN_OVERLAP`` governs branch 3's overlap term. Branches 2 and 3
    are deliberately not mutually simplified: the constants are separate
    re-calibration controls, and the gap between them is the ambiguous band.
    """
    speech_windows = [row for row in window_stats if row.speech_frames > 0]
    if not speech_windows:
        return SpeakerEvidenceDecision(
            speaker_evidence="none",
            multi_window_fraction=0.0,
            mean_window_overlap_share=0.0,
        )

    multi_window_count = sum(1 for row in speech_windows if row.active_slot_count > 1)
    multi_window_fraction = multi_window_count / len(speech_windows)
    mean_window_overlap_share = sum(
        row.overlap_frames / row.speech_frames for row in speech_windows
    ) / len(speech_windows)

    if (
        multi_window_fraction >= SPEAKER_EVIDENCE_MULTI_MIN
        or overlap_fraction >= DIARIZE_MIN_OVERLAP
    ):
        speaker_evidence = "multi"
    elif (
        multi_window_fraction < SPEAKER_EVIDENCE_SINGLE_MAX
        and mean_window_overlap_share < DIARIZE_MIN_OVERLAP
    ):
        speaker_evidence = "single"
    else:
        speaker_evidence = "multi"

    return SpeakerEvidenceDecision(
        speaker_evidence=speaker_evidence,
        multi_window_fraction=multi_window_fraction,
        mean_window_overlap_share=mean_window_overlap_share,
    )


def _get_overlap_session() -> ort.InferenceSession:
    """Return a cached ONNX InferenceSession for the pyannote overlap model."""
    global _overlap_session
    import onnxruntime as ort

    if _overlap_session is None:
        from solstone.observe.model_assets import resolve_pyannote_segmentation_model
        from solstone.observe.transcribe.main import _select_onnx_providers

        pyannote_model_path = resolve_pyannote_segmentation_model()
        if not pyannote_model_path.is_file():
            raise FileNotFoundError(
                f"pyannote model not found at {pyannote_model_path}. "
                "Run `make install` to verify the bundled asset."
            )

        providers = _select_onnx_providers()
        start = time.monotonic()
        _overlap_session = ort.InferenceSession(
            str(pyannote_model_path),
            providers=providers,
        )
        elapsed_ms = (time.monotonic() - start) * 1000.0
        logger.info(
            "pyannote overlap session loaded (providers=%s, elapsed=%.1fms)",
            _overlap_session.get_providers(),
            elapsed_ms,
        )

    return _overlap_session


def compute_overlap_fraction(
    audio: np.ndarray, sample_rate: int = SAMPLE_RATE
) -> float:
    """Compute the speech-conditioned overlap fraction for an audio segment."""
    import numpy as np

    if sample_rate != SAMPLE_RATE:
        raise ValueError(
            f"pyannote overlap detector requires {SAMPLE_RATE} Hz audio, got {sample_rate}"
        )

    session = _get_overlap_session()
    input_name = session.get_inputs()[0].name

    window_samples = WINDOW_S * sample_rate
    stride_samples = STRIDE_S * sample_rate

    audio_f32 = np.asarray(audio, dtype=np.float32)
    if len(audio_f32) < window_samples:
        pad = window_samples - len(audio_f32)
        audio_padded = np.concatenate([audio_f32, np.zeros(pad, dtype=np.float32)])
    else:
        audio_padded = audio_f32

    starts = list(range(0, len(audio_padded) - window_samples + 1, stride_samples))
    final_start = max(0, len(audio_padded) - window_samples)
    if not starts:
        starts = [final_start]
    elif starts[-1] != final_start:
        starts.append(final_start)

    samples_per_frame = window_samples / FRAMES_PER_WINDOW
    num_frames = int(np.ceil(len(audio_padded) / samples_per_frame))
    accum = np.zeros((num_frames, 7), dtype=np.float64)
    counts = np.zeros((num_frames,), dtype=np.int32)

    for start_sample in starts:
        chunk = audio_padded[start_sample : start_sample + window_samples][
            None, None, :
        ]
        log_probs = session.run(None, {input_name: chunk})[0][0]
        frame_start = int(round(start_sample / samples_per_frame))
        frame_end = frame_start + log_probs.shape[0]
        if frame_end > num_frames:
            frame_end = num_frames
            log_probs = log_probs[: frame_end - frame_start]
        accum[frame_start:frame_end] += log_probs.astype(np.float64)
        counts[frame_start:frame_end] += 1

    counts = np.maximum(counts, 1)
    avg_log_probs = (accum / counts[:, None]).astype(np.float32)
    argmax = avg_log_probs.argmax(axis=-1)
    speech_count = int((argmax >= 1).sum())
    if speech_count == 0:
        return 0.0

    overlap_count = int(np.isin(argmax, OVERLAP_CLASSES).sum())
    return float(overlap_count / speech_count)


def compute_overlap_and_logprobs(
    audio: np.ndarray, sample_rate: int = SAMPLE_RATE
) -> OverlapInferenceResult:
    """Compute overlap fraction and return the pyannote log-probs for reuse.

    Uses a 2-second stride (vs. the 5-second stride in compute_overlap_fraction)
    for better per-frame resolution.  The returned avg_log_probs array has shape
    (num_frames, 7) and can be passed directly to diarize() to skip a second
    pyannote inference pass.
    """
    import numpy as np

    if sample_rate != SAMPLE_RATE:
        raise ValueError(
            f"pyannote overlap detector requires {SAMPLE_RATE} Hz audio, got {sample_rate}"
        )

    session = _get_overlap_session()
    input_name = session.get_inputs()[0].name

    window_samples = WINDOW_S * sample_rate
    stride_samples = _DIARIZE_STRIDE_S * sample_rate

    audio_f32 = np.asarray(audio, dtype=np.float32)
    if len(audio_f32) < window_samples:
        pad = window_samples - len(audio_f32)
        audio_padded = np.concatenate([audio_f32, np.zeros(pad, dtype=np.float32)])
    else:
        audio_padded = audio_f32

    starts = list(range(0, len(audio_padded) - window_samples + 1, stride_samples))
    final_start = max(0, len(audio_padded) - window_samples)
    if not starts:
        starts = [final_start]
    elif starts[-1] != final_start:
        starts.append(final_start)

    samples_per_frame = window_samples / FRAMES_PER_WINDOW
    num_frames = int(np.ceil(len(audio_padded) / samples_per_frame))
    accum = np.zeros((num_frames, 7), dtype=np.float64)
    counts = np.zeros((num_frames,), dtype=np.int32)
    window_stats: list[SpeakerWindowStats] = []

    for start_sample in starts:
        chunk = audio_padded[start_sample : start_sample + window_samples][
            None, None, :
        ]
        log_probs = session.run(None, {input_name: chunk})[0][0]
        frame_start = int(round(start_sample / samples_per_frame))
        frame_end = frame_start + log_probs.shape[0]
        if frame_end > num_frames:
            frame_end = num_frames
            log_probs = log_probs[: frame_end - frame_start]
        window_stats.append(_speaker_window_stats(log_probs))
        accum[frame_start:frame_end] += log_probs.astype(np.float64)
        counts[frame_start:frame_end] += 1

    counts = np.maximum(counts, 1)
    avg_log_probs = (accum / counts[:, None]).astype(np.float32)

    argmax = avg_log_probs.argmax(axis=-1)
    speech_count = int((argmax >= 1).sum())
    if speech_count == 0:
        return OverlapInferenceResult(0.0, avg_log_probs, tuple(window_stats))

    overlap_count = int(np.isin(argmax, OVERLAP_CLASSES).sum())
    return OverlapInferenceResult(
        float(overlap_count / speech_count),
        avg_log_probs,
        tuple(window_stats),
    )


def compute_overlap_fraction_for_wav(path: Path) -> float:
    """Load a WAV file and compute its speech-conditioned overlap fraction."""
    import numpy as np
    import soundfile as sf
    from scipy.signal import resample_poly

    audio, sample_rate = sf.read(path, dtype="float32", always_2d=False)
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    if sample_rate != SAMPLE_RATE:
        audio = resample_poly(audio, SAMPLE_RATE, sample_rate).astype(np.float32)
        sample_rate = SAMPLE_RATE
    return compute_overlap_fraction(np.asarray(audio, dtype=np.float32), sample_rate)
