src/turn_detector/features.py

"""
Classical acoustic feature extraction for the energy/silence baseline
(EXP-001) and the lightweight classifier baseline (EXP-002/EXP-003).

Implemented with numpy + scipy only (no librosa dependency) so this module
can be developed and unit-tested in this sandbox, which has numpy/scipy but
not librosa/soundfile/torch available and no network access to install
them. If librosa is available in the actual training environment it is a
fine drop-in replacement for `mfcc()` — the manual mel-filterbank + DCT
implementation here exists specifically so this module doesn't have a hard
librosa dependency, not because librosa's MFCC is worse.

All functions operate on a 1D float32 numpy array of mono audio samples at
a known sample rate. Nothing in this file reads audio from disk or the
network — see audio_io.py for that.
"""

from __future__ import annotations

from dataclasses import dataclass, asdict

import numpy as np
from scipy.fftpack import dct

EPS = 1e-10


# ---------------------------------------------------------------------------
# Frame-level primitives
# ---------------------------------------------------------------------------

def frame_signal(audio: np.ndarray, frame_len: int, hop_len: int) -> np.ndarray:
    """Split audio into overlapping frames. Returns shape (n_frames, frame_len).
    Frames past the end are dropped (no zero-padding) — for feature
    extraction we don't want to invent energy from padding.
    """
    n = len(audio)
    if n < frame_len:
        return np.empty((0, frame_len), dtype=audio.dtype)
    n_frames = 1 + (n - frame_len) // hop_len
    idx = np.arange(frame_len)[None, :] + hop_len * np.arange(n_frames)[:, None]
    return audio[idx]


def rms_per_frame(frames: np.ndarray) -> np.ndarray:
    return np.sqrt(np.mean(frames.astype(np.float64) ** 2, axis=1) + EPS)


def zcr_per_frame(frames: np.ndarray) -> np.ndarray:
    signs = np.sign(frames)
    signs[signs == 0] = 1
    return np.mean(np.abs(np.diff(signs, axis=1)) > 0, axis=1)


# ---------------------------------------------------------------------------
# Spectral features
# ---------------------------------------------------------------------------

def _power_spectrum(frames: np.ndarray, sr: int) -> tuple[np.ndarray, np.ndarray]:
    n_fft = frames.shape[1]
    window = np.hanning(n_fft)
    windowed = frames * window[None, :]
    spec = np.fft.rfft(windowed, axis=1)
    power = np.abs(spec) ** 2
    freqs = np.fft.rfftfreq(n_fft, d=1.0 / sr)
    return power, freqs


def spectral_centroid(frames: np.ndarray, sr: int) -> np.ndarray:
    power, freqs = _power_spectrum(frames, sr)
    total = power.sum(axis=1) + EPS
    return (power * freqs[None, :]).sum(axis=1) / total


def spectral_bandwidth(frames: np.ndarray, sr: int) -> np.ndarray:
    power, freqs = _power_spectrum(frames, sr)
    total = power.sum(axis=1) + EPS
    centroid = (power * freqs[None, :]).sum(axis=1) / total
    diff_sq = (freqs[None, :] - centroid[:, None]) ** 2
    return np.sqrt((power * diff_sq).sum(axis=1) / total)


def spectral_rolloff(frames: np.ndarray, sr: int, roll_percent: float = 0.85) -> np.ndarray:
    power, freqs = _power_spectrum(frames, sr)
    cumsum = np.cumsum(power, axis=1)
    total = cumsum[:, -1:] + EPS
    threshold = roll_percent * total
    idx = (cumsum >= threshold).argmax(axis=1)
    return freqs[idx]


def mel_filterbank(sr: int, n_fft: int, n_mels: int = 26, fmin: float = 0.0, fmax: float | None = None) -> np.ndarray:
    """Standard triangular mel filterbank, computed with numpy only."""
    fmax = fmax or sr / 2

    def hz_to_mel(f):
        return 2595.0 * np.log10(1.0 + f / 700.0)

    def mel_to_hz(m):
        return 700.0 * (10 ** (m / 2595.0) - 1.0)

    mel_min, mel_max = hz_to_mel(fmin), hz_to_mel(fmax)
    mel_points = np.linspace(mel_min, mel_max, n_mels + 2)
    hz_points = mel_to_hz(mel_points)
    bin_points = np.floor((n_fft + 1) * hz_points / sr).astype(int)

    fbank = np.zeros((n_mels, n_fft // 2 + 1))
    for m in range(1, n_mels + 1):
        left, center, right = bin_points[m - 1], bin_points[m], bin_points[m + 1]
        if center == left:
            center += 1
        if right == center:
            right += 1
        for k in range(left, min(center, fbank.shape[1])):
            fbank[m - 1, k] = (k - left) / (center - left)
        for k in range(center, min(right, fbank.shape[1])):
            fbank[m - 1, k] = (right - k) / (right - center)
    return fbank


def mfcc(frames: np.ndarray, sr: int, n_mfcc: int = 13, n_mels: int = 26) -> np.ndarray:
    """Return shape (n_frames, n_mfcc). Manual implementation — see module
    docstring for why this doesn't depend on librosa.
    """
    power, _ = _power_spectrum(frames, sr)
    fbank = mel_filterbank(sr, frames.shape[1], n_mels=n_mels)
    mel_energy = power @ fbank.T
    log_mel = np.log(mel_energy + EPS)
    coeffs = dct(log_mel, type=2, axis=1, norm="ortho")[:, :n_mfcc]
    return coeffs


# ---------------------------------------------------------------------------
# Silence / energy-trend features (used directly by the EXP-001 baseline)
# ---------------------------------------------------------------------------

def trailing_silence_duration(
    audio: np.ndarray,
    sr: int,
    frame_len_ms: float = 20.0,
    hop_len_ms: float = 10.0,
    silence_rms_threshold: float = 0.01,
) -> float:
    """Duration in seconds of trailing silence at the end of the clip,
    using a fixed RMS threshold. This is the "classic VAD-style endpointing"
    signal used by EXP-001. The threshold is a tunable hyperparameter —
    see scripts for the tuning procedure (must be tuned on dev/val split
    only, never on the held-out test set, per project rules).
    """
    frame_len = max(1, int(sr * frame_len_ms / 1000))
    hop_len = max(1, int(sr * hop_len_ms / 1000))
    frames = frame_signal(audio, frame_len, hop_len)
    if len(frames) == 0:
        return len(audio) / sr  # whole clip shorter than one frame — treat as silence-length itself
    rms = rms_per_frame(frames)
    is_silent = rms < silence_rms_threshold
    # count trailing silent frames
    n_trailing = 0
    for v in is_silent[::-1]:
        if v:
            n_trailing += 1
        else:
            break
    return n_trailing * hop_len_ms / 1000.0


def energy_slope(rms: np.ndarray, n_recent: int = 5) -> float:
    """Linear-fit slope of RMS energy over the last `n_recent` frames.
    Negative slope = energy decaying (consistent with trailing off).
    Returns 0.0 if too few frames.
    """
    if len(rms) < 2:
        return 0.0
    recent = rms[-n_recent:]
    if len(recent) < 2:
        return 0.0
    x = np.arange(len(recent))
    slope, _ = np.polyfit(x, recent, 1)
    return float(slope)


# ---------------------------------------------------------------------------
# Full feature vector for the classifier baseline (EXP-002 / EXP-003)
# ---------------------------------------------------------------------------

@dataclass
class FeatureConfig:
    sr: int = 16_000
    frame_len_ms: float = 25.0
    hop_len_ms: float = 10.0
    n_mfcc: int = 13
    n_mels: int = 26
    silence_rms_threshold: float = 0.01
    # Recent-window sizes for temporal features (EXP-003 in this phase's
    # experiment list — "is the tail of the audio more predictive than the
    # whole-clip statistics?"). Windows shorter than one frame or longer
    # than the clip are silently skipped per-clip (documented, not hidden).
    recent_windows_ms: tuple = (100, 250, 500, 750, 1000)


def _summary_stats(x: np.ndarray, prefix: str) -> dict:
    if len(x) == 0:
        return {f"{prefix}_mean": 0.0, f"{prefix}_std": 0.0, f"{prefix}_min": 0.0, f"{prefix}_max": 0.0}
    return {
        f"{prefix}_mean": float(np.mean(x)),
        f"{prefix}_std": float(np.std(x)),
        f"{prefix}_min": float(np.min(x)),
        f"{prefix}_max": float(np.max(x)),
    }


def extract_global_features(audio: np.ndarray, cfg: FeatureConfig = FeatureConfig()) -> dict:
    """Features computed over the entire clip — the "global statistics"
    half of EXP-002/EXP-003's comparison against recent-window features.
    """
    sr = cfg.sr
    frame_len = max(1, int(sr * cfg.frame_len_ms / 1000))
    hop_len = max(1, int(sr * cfg.hop_len_ms / 1000))
    frames = frame_signal(audio, frame_len, hop_len)

    feats: dict = {"duration_sec": len(audio) / sr}

    if len(frames) == 0:
        # Clip shorter than one frame: return zeroed features rather than
        # crashing, but flag it so downstream code can filter/inspect these.
        feats["too_short_for_framing"] = True
        return feats

    feats["too_short_for_framing"] = False

    rms = rms_per_frame(frames)
    zcr = zcr_per_frame(frames)
    centroid = spectral_centroid(frames, sr)
    bandwidth = spectral_bandwidth(frames, sr)
    rolloff = spectral_rolloff(frames, sr)
    coeffs = mfcc(frames, sr, n_mfcc=cfg.n_mfcc, n_mels=cfg.n_mels)

    feats.update(_summary_stats(rms, "rms"))
    feats.update(_summary_stats(zcr, "zcr"))
    feats.update(_summary_stats(centroid, "centroid"))
    feats.update(_summary_stats(bandwidth, "bandwidth"))
    feats.update(_summary_stats(rolloff, "rolloff"))
    for i in range(coeffs.shape[1]):
        feats.update(_summary_stats(coeffs[:, i], f"mfcc{i}"))

    feats["energy_slope"] = energy_slope(rms)
    feats["trailing_silence_sec"] = trailing_silence_duration(
        audio, sr, cfg.frame_len_ms, cfg.hop_len_ms, cfg.silence_rms_threshold
    )
    return feats


def extract_recent_window_features(audio: np.ndarray, cfg: FeatureConfig = FeatureConfig()) -> dict:
    """Features computed over only the last N ms of audio, for each window
    size in cfg.recent_windows_ms. This directly supports the Phase 2
    research question: does the tail of the clip predict the label better
    than whole-clip statistics? Windows longer than the available audio are
    skipped (recorded as NaN) rather than padded with fabricated silence,
    so downstream analysis can see exactly which clips had enough audio for
    each window size.
    """
    sr = cfg.sr
    feats: dict = {}
    for w_ms in cfg.recent_windows_ms:
        w_samples = int(sr * w_ms / 1000)
        key_prefix = f"tail{w_ms}ms"
        if len(audio) < w_samples:
            feats[f"{key_prefix}_available"] = False
            feats[f"{key_prefix}_rms_mean"] = float("nan")
            continue
        window = audio[-w_samples:]
        frame_len = max(1, int(sr * cfg.frame_len_ms / 1000))
        hop_len = max(1, int(sr * cfg.hop_len_ms / 1000))
        frames = frame_signal(window, min(frame_len, w_samples), max(1, min(hop_len, w_samples)))
        feats[f"{key_prefix}_available"] = True
        if len(frames) == 0:
            rms_val = rms_per_frame(window.reshape(1, -1))
        else:
            rms_val = rms_per_frame(frames)
        feats[f"{key_prefix}_rms_mean"] = float(np.mean(rms_val))
        feats[f"{key_prefix}_rms_slope"] = energy_slope(rms_val, n_recent=len(rms_val))
    return feats


def extract_features(audio: np.ndarray, cfg: FeatureConfig = FeatureConfig(), include_recent_windows: bool = True) -> dict:
    feats = extract_global_features(audio, cfg)
    if include_recent_windows:
        feats.update(extract_recent_window_features(audio, cfg))
    return feats