src/turn_detection/baselines.py

"""Interpretable baselines for endpoint classification and pause control.

These baselines are deliberately small enough to audit.  The acoustic model
uses waveform-derived statistics only; it never consumes source, language, or
synthetic flags that could become dataset shortcuts.
"""

from __future__ import annotations

import io
import json
import math
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from turn_detection.runtime.replay import PauseCheckpoint, ReplayRecord, summarize_replay
from turn_detection.runtime.types import TurnDecision, TurnState

AUDIO_FEATURE_NAMES = (
    "log_duration_seconds",
    "rms_db",
    "tail_rms_db",
    "tail_energy_delta_db",
    "tail_silence_fraction",
    "tail_zero_crossing_rate",
    "spectral_centroid_nyquist",
    "spectral_rolloff85_nyquist",
    "periodicity",
    "energy_slope",
)


def _numpy() -> Any:
    try:
        import numpy as np
    except ImportError as exc:  # pragma: no cover - base dependency
        raise RuntimeError("audio baselines require numpy") from exc
    return np


def _sigmoid(values: Any) -> Any:
    np = _numpy()
    clipped = np.clip(values, -30.0, 30.0)
    return 1.0 / (1.0 + np.exp(-clipped))


def decode_audio_value(value: Any) -> tuple[Any, int]:
    """Decode a Hugging Face audio value without materializing a dataset."""

    np = _numpy()
    try:
        import soundfile as sf
    except ImportError as exc:  # pragma: no cover - base dependency
        raise RuntimeError("audio baselines require soundfile") from exc

    declared_rate: int | None = None
    source: Any = value
    if isinstance(value, Mapping):
        declared_rate = value.get("sampling_rate", value.get("sample_rate"))
        if value.get("array") is not None:
            array = np.asarray(value["array"])
            source_rate = int(declared_rate) if declared_rate else 16_000
            return _mono_float(array), source_rate
        if value.get("samples") is not None:
            array = np.asarray(value["samples"])
            source_rate = int(declared_rate) if declared_rate else 16_000
            return _mono_float(array), source_rate
        if value.get("bytes") is not None:
            source = io.BytesIO(value["bytes"])
        elif value.get("path"):
            source = value["path"]
        else:
            raise ValueError("audio mapping contains no bytes, path, array, or samples")
    elif isinstance(value, (bytes, bytearray)):
        source = io.BytesIO(bytes(value))
    elif isinstance(value, (str, Path)):
        source = str(value)
    else:
        return _mono_float(np.asarray(value)), int(declared_rate or 16_000)

    array, decoded_rate = sf.read(source, dtype="float32", always_2d=True)
    return _mono_float(array), int(declared_rate or decoded_rate)


def _mono_float(audio: Any) -> Any:
    np = _numpy()
    samples = np.asarray(audio)
    original_dtype = samples.dtype
    if samples.size == 0:
        raise ValueError("audio cannot be empty")
    if samples.ndim == 2:
        channel_axis = 1 if samples.shape[1] <= 8 else 0
        samples = samples.astype(np.float32).mean(axis=channel_axis)
    elif samples.ndim != 1:
        raise ValueError(f"expected one- or two-dimensional audio, got {samples.shape}")
    if np.issubdtype(original_dtype, np.integer):
        info = np.iinfo(original_dtype)
        samples = samples.astype(np.float32) / float(max(abs(info.min), info.max))
    else:
        samples = samples.astype(np.float32, copy=False)
    samples = np.nan_to_num(samples, nan=0.0, posinf=1.0, neginf=-1.0)
    peak = float(np.max(np.abs(samples)))
    if peak > 1.0:
        samples = samples / peak
    return np.clip(samples, -1.0, 1.0)


def _linear_resample(samples: Any, source_rate: int, target_rate: int) -> Any:
    np = _numpy()
    if source_rate == target_rate:
        return samples.astype(np.float32, copy=False)
    output_length = max(1, round(len(samples) * target_rate / source_rate))
    old_positions = np.arange(len(samples), dtype=np.float64)
    new_positions = np.linspace(0.0, max(0, len(samples) - 1), output_length)
    return np.interp(new_positions, old_positions, samples).astype(np.float32)


def extract_audio_statistics(
    audio: Any,
    sample_rate: int,
    *,
    target_rate: int = 16_000,
    max_seconds: float = 4.0,
) -> Any:
    """Return ten finite, interpretable suffix-acoustic features."""

    np = _numpy()
    if sample_rate <= 0 or target_rate <= 0 or max_seconds <= 0:
        raise ValueError("sample rates and max_seconds must be positive")
    samples = _linear_resample(_mono_float(audio), sample_rate, target_rate)
    duration_seconds = len(samples) / target_rate
    samples = samples[-round(max_seconds * target_rate) :]
    tail = samples[-min(len(samples), round(0.40 * target_rate)) :]
    body = samples[: -len(tail)] if len(samples) > len(tail) else samples
    eps = 1e-10

    def rms_db(values: Any) -> float:
        return 10.0 * math.log10(float(np.mean(values * values)) + eps)

    total_db = rms_db(samples)
    tail_db = rms_db(tail)
    body_db = rms_db(body)
    tail_silence_fraction = float(np.mean(np.abs(tail) < 0.01))
    tail_zcr = float(np.mean((tail[1:] >= 0) != (tail[:-1] >= 0))) if len(tail) > 1 else 0.0

    analysis = tail[-min(len(tail), target_rate) :]
    windowed = analysis * np.hanning(len(analysis)).astype(np.float32)
    spectrum = np.abs(np.fft.rfft(windowed)) ** 2
    frequencies = np.fft.rfftfreq(len(windowed), d=1.0 / target_rate)
    spectral_sum = float(spectrum.sum()) + eps
    centroid = float((spectrum * frequencies).sum() / spectral_sum) / (target_rate / 2.0)
    cumulative = np.cumsum(spectrum)
    rolloff_index = int(np.searchsorted(cumulative, 0.85 * cumulative[-1]))
    rolloff_index = min(rolloff_index, len(frequencies) - 1)
    rolloff = float(frequencies[rolloff_index]) / (target_rate / 2.0)

    centered = analysis - float(analysis.mean())
    energy = float(np.dot(centered, centered)) + eps
    minimum_lag = max(1, target_rate // 400)
    maximum_lag = min(len(centered) - 1, target_rate // 60)
    periodicity = 0.0
    if maximum_lag >= minimum_lag and energy > eps:
        correlations = [
            float(np.dot(centered[:-lag], centered[lag:])) / energy
            for lag in range(minimum_lag, maximum_lag + 1)
        ]
        periodicity = max(0.0, max(correlations, default=0.0))

    frame_length = max(1, round(0.10 * target_rate))
    frame_hop = max(1, round(0.05 * target_rate))
    frame_rms: list[float] = []
    for start in range(0, max(1, len(samples) - frame_length + 1), frame_hop):
        frame = samples[start : start + frame_length]
        if len(frame):
            frame_rms.append(rms_db(frame))
    if len(frame_rms) >= 2:
        x = np.linspace(-1.0, 1.0, len(frame_rms))
        energy_slope = float(np.dot(x, np.asarray(frame_rms) - np.mean(frame_rms))) / float(
            np.dot(x, x)
        )
    else:
        energy_slope = 0.0

    features = np.asarray(
        [
            math.log1p(duration_seconds),
            total_db,
            tail_db,
            tail_db - body_db,
            tail_silence_fraction,
            tail_zcr,
            centroid,
            rolloff,
            periodicity,
            energy_slope,
        ],
        dtype=np.float64,
    )
    if not bool(np.isfinite(features).all()):
        raise ValueError("audio statistics contained non-finite values")
    return features


@dataclass(frozen=True, slots=True)
class LogisticBaseline:
    feature_names: tuple[str, ...]
    mean: tuple[float, ...]
    scale: tuple[float, ...]
    weights: tuple[float, ...]
    bias: float
    threshold: float = 0.5
    sample_rate: int = 16_000
    max_seconds: float = 4.0

    def __post_init__(self) -> None:
        lengths = {len(self.feature_names), len(self.mean), len(self.scale), len(self.weights)}
        if len(lengths) != 1:
            raise ValueError("model vectors must have the same length")
        if any(value <= 0 for value in self.scale):
            raise ValueError("standardization scales must be positive")
        if not 0.0 <= self.threshold <= 1.0:
            raise ValueError("threshold must be in [0, 1]")

    def predict_proba(self, features: Any) -> Any:
        np = _numpy()
        matrix = np.asarray(features, dtype=np.float64)
        standardized = (matrix - np.asarray(self.mean)) / np.asarray(self.scale)
        return _sigmoid(standardized @ np.asarray(self.weights) + self.bias)

    def to_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        payload["model_type"] = "audio_statistics_logistic_regression"
        payload["feature_names"] = list(self.feature_names)
        payload["mean"] = list(self.mean)
        payload["scale"] = list(self.scale)
        payload["weights"] = list(self.weights)
        return payload

    @classmethod
    def from_dict(cls, payload: Mapping[str, Any]) -> LogisticBaseline:
        values = dict(payload)
        values.pop("model_type", None)
        for name in ("feature_names", "mean", "scale", "weights"):
            values[name] = tuple(values[name])
        return cls(**values)


def fit_logistic_baseline(
    features: Any,
    labels: Any,
    *,
    epochs: int = 800,
    learning_rate: float = 0.05,
    l2: float = 1e-3,
    feature_names: Sequence[str] | None = None,
) -> LogisticBaseline:
    """Fit deterministic class-balanced logistic regression with full-batch GD."""

    np = _numpy()
    matrix = np.asarray(features, dtype=np.float64)
    targets = np.asarray(labels, dtype=np.float64).reshape(-1)
    if matrix.ndim != 2 or matrix.shape[0] != len(targets) or matrix.shape[1] == 0:
        raise ValueError("features must be [examples, dimensions] and align with labels")
    if len(targets) == 0 or not bool(np.isin(targets, [0.0, 1.0]).all()):
        raise ValueError("labels must be a non-empty binary vector")
    resolved_names = tuple(feature_names or ())
    if not resolved_names:
        resolved_names = (
            AUDIO_FEATURE_NAMES
            if matrix.shape[1] == len(AUDIO_FEATURE_NAMES)
            else tuple(f"feature_{index}" for index in range(matrix.shape[1]))
        )
    if len(resolved_names) != matrix.shape[1] or len(set(resolved_names)) != len(resolved_names):
        raise ValueError("feature_names must be unique and match the feature dimension")
    positives = float(targets.sum())
    negatives = float(len(targets) - positives)
    if positives == 0 or negatives == 0:
        raise ValueError("both endpoint classes are required")
    if epochs < 1 or learning_rate <= 0 or l2 < 0:
        raise ValueError("epochs/learning_rate must be positive and l2 non-negative")

    mean = matrix.mean(axis=0)
    scale = matrix.std(axis=0)
    scale = np.where(scale < 1e-8, 1.0, scale)
    standardized = (matrix - mean) / scale
    weights = np.zeros(matrix.shape[1], dtype=np.float64)
    bias = 0.0
    example_weights = np.where(
        targets == 1.0, len(targets) / (2 * positives), len(targets) / (2 * negatives)
    )
    denominator = float(example_weights.sum())
    for _ in range(epochs):
        probabilities = _sigmoid(standardized @ weights + bias)
        residual = (probabilities - targets) * example_weights
        weights -= learning_rate * ((standardized.T @ residual) / denominator + l2 * weights)
        bias -= learning_rate * float(residual.sum() / denominator)

    return LogisticBaseline(
        feature_names=resolved_names,
        mean=tuple(float(value) for value in mean),
        scale=tuple(float(value) for value in scale),
        weights=tuple(float(value) for value in weights),
        bias=float(bias),
    )


def extract_manifest_features(
    rows: Iterable[Mapping[str, Any]],
    *,
    source_root: str | Path,
    max_examples: int | None = None,
    sample_rate: int = 16_000,
    max_seconds: float = 4.0,
) -> tuple[Any, Any, list[dict[str, Any]]]:
    """Resolve manifest audio lazily and produce a numeric matrix plus provenance."""

    from turn_detection.data import ManifestRecordResolver

    np = _numpy()
    if max_examples is not None and max_examples < 1:
        raise ValueError("max_examples must be positive when provided")
    resolver = ManifestRecordResolver(source_root=source_root, max_cached_row_groups=1)
    vectors: list[Any] = []
    labels: list[int] = []
    provenance: list[dict[str, Any]] = []
    try:
        for row in rows:
            if max_examples is not None and len(vectors) >= max_examples:
                break
            endpoint = row.get("endpoint", row.get("endpoint_bool"))
            if endpoint not in (False, True, 0, 1):
                raise ValueError(f"manifest row {row.get('record_id')!r} has no binary endpoint")
            source_record = resolver.resolve(row, columns=("audio",))
            audio, decoded_rate = decode_audio_value(source_record["audio"])
            vectors.append(
                extract_audio_statistics(
                    audio,
                    decoded_rate,
                    target_rate=sample_rate,
                    max_seconds=max_seconds,
                )
            )
            labels.append(int(bool(endpoint)))
            provenance.append(
                {
                    "record_id": str(row.get("record_id", "")),
                    "group_id": row.get("group_id"),
                    "split": row.get("split"),
                    "label": int(bool(endpoint)),
                }
            )
    finally:
        resolver.clear()
    if not vectors:
        raise ValueError("manifest selection produced no examples")
    return np.stack(vectors), np.asarray(labels, dtype=np.int64), provenance


def fixed_timeout_sweep(
    checkpoints: Iterable[PauseCheckpoint],
    timeouts_ms: Sequence[float],
) -> list[dict[str, Any]]:
    """Evaluate model-free policies that fire after a fixed VAD silence."""

    materialized = list(checkpoints)
    results: list[dict[str, Any]] = []
    for timeout in timeouts_ms:
        if not math.isfinite(timeout) or timeout < 0:
            raise ValueError("timeouts must be finite and non-negative")
        replay_records: list[ReplayRecord] = []
        ended_turns: set[str] = set()
        for checkpoint in materialized:
            should_end = checkpoint.silence_ms >= timeout
            emit_response = should_end and checkpoint.turn_id not in ended_turns
            if emit_response:
                ended_turns.add(checkpoint.turn_id)
            state = TurnState.END if checkpoint.turn_id in ended_turns else TurnState.HOLD
            replay_records.append(
                ReplayRecord(
                    checkpoint=checkpoint,
                    decision=TurnDecision(
                        state=state,
                        endpoint_probability=None,
                        threshold=None,
                        silence_ms=checkpoint.silence_ms,
                        reason=f"fixed_timeout_{timeout:g}ms",
                        timestamp_ms=checkpoint.timestamp_ms,
                        emit_response=emit_response,
                    ),
                )
            )
        results.append({"timeout_ms": float(timeout), **summarize_replay(replay_records)})
    return results


def load_checkpoints_jsonl(path: str | Path) -> list[PauseCheckpoint]:
    checkpoints: list[PauseCheckpoint] = []
    with Path(path).open(encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            try:
                value = json.loads(line)
                checkpoints.append(PauseCheckpoint(**value))
            except (json.JSONDecodeError, TypeError, ValueError) as exc:
                raise ValueError(f"invalid checkpoint at {path}:{line_number}") from exc
    return checkpoints