speaker_processing.py

"""Core audio processing for Voice Studio.

Pipeline:
    1. extract_audio / load_audio  -> extract a mono WAV from a video or audio file (ffmpeg)
    2. diarize                     -> split speech into one slot per voice actor
                                      (pyannote-audio 3.1 when an HF token is
                                      available, else ECAPA-TDNN window
                                      classification with overlap detection;
                                      MFCC + KMeans as last resort)
    3. profile_speakers            -> estimate gender & age per actor from pitch statistics
    4. auto_gains / mix            -> apply per-actor volume gains (manual or auto-balanced)

The number of speakers is detected automatically from embedding similarity, so the
caller never has to know how many people are in the recording. Automatic gender/age
estimates remain lightweight, pitch-based heuristics meant to pre-fill the interface;
they can (and should) be overridden manually in the UI.
"""

import io
import os
import subprocess
import tempfile
import threading
import time

import numpy as np
import librosa
import soundfile as sf

try:  # scikit-learn is optional: without it, all speech goes to a single actor
    from sklearn.cluster import KMeans
    from sklearn.metrics import silhouette_score
    _HAS_SKLEARN = True
except ImportError:  # pragma: no cover
    _HAS_SKLEARN = False

SAMPLE_RATE = 22050
FRAME_MS = 30
HOP_MS = 15
ENERGY_DB_THRESH = -40.0
MIN_SILENCE_S = 0.4
MIN_SPEECH_S = 0.5
MIN_SEGMENT_S = 0.5  # segments shorter than this are not used for clustering

# Optional speechbrain / torch (ECAPA-TDNN embeddings). Imported lazily so that
# starting the app doesn't pay torch's import cost; ``None`` means "not probed".
_HAS_SPEECHBRAIN = None
_TORCH = None
_ENCODER_CLASSIFIER = None

ECAPA_SOURCE = "speechbrain/spkrec-ecapa-voxceleb"
ECAPA_SAVEDIR = os.path.join(tempfile.gettempdir(), "voice_studio_ecapa")
EMBEDDING_SR = 16000          # ECAPA-TDNN is trained on 16 kHz audio
EMBEDDING_MIN_S = 1.0         # segments shorter than this are not used for clustering
ASSIGN_MIN_S = 0.3            # shortest window still embedded (to assign it to a speaker)
MAX_SPEAKERS = 10
WINDOW_S = 1.0                # classification window length (overlap detection granularity)
WINDOW_HOP_S = 0.5            # window stride (overlapping tiles -> finer boundaries)
OVERLAP_GAP = 0.10            # top-2 centroid similarity gap -> window counts as "both speakers"

# Optional audeering wav2vec2 age & gender model (ONNX, loaded via audonnx).
# ``None`` = not probed, ``False`` = unavailable (fall back to pitch heuristics).
_AGE_GENDER_MODEL = None
_AGE_GENDER_LOCK = threading.Lock()
AGE_GENDER_URL = ("https://zenodo.org/record/7761387/files/"
                  "w2v2-L-robust-6-age-gender.25c844af-1.1.1.zip")
AGE_GENDER_SAVEDIR = os.path.join(tempfile.gettempdir(), "voice_studio_age_gender")
AGE_GENDER_SR = 16000
AGE_GENDER_LABELS = ["Female", "Male", "Child"]


# --------------------------------------------------------------------------- #
# 1. Audio extraction
# --------------------------------------------------------------------------- #
def _ffmpeg_to_wav(source_path: str, out_wav: str, target_sr: int = SAMPLE_RATE) -> None:
    """Convert any ffmpeg-supported file to a mono WAV at ``target_sr``."""
    cmd = [
        "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
        "-i", source_path,
        "-ac", "1", "-ar", str(target_sr),
        "-vn", out_wav,
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        raise RuntimeError(
            "ffmpeg failed to read this file. Is ffmpeg installed?\n"
            + (proc.stderr or "").strip()
        )


def load_audio(source_path: str, target_sr: int = SAMPLE_RATE):
    """Load a video/audio file as a mono float array ``(y, sr)`` in [-1, 1]."""
    fd, tmp_wav = tempfile.mkstemp(suffix=".wav")
    os.close(fd)
    try:
        _ffmpeg_to_wav(source_path, tmp_wav, target_sr)
        return librosa.load(tmp_wav, sr=target_sr, mono=True)
    finally:
        if os.path.exists(tmp_wav):
            os.remove(tmp_wav)


# --------------------------------------------------------------------------- #
# 2. Diarization (VAD + speaker clustering)
# --------------------------------------------------------------------------- #
_VAD_MODEL = None
_VAD_LOCK = threading.Lock()
VAD_SR = 16000


def _get_silero_vad():
    """Lazily load the Silero VAD model; ``False`` when unavailable."""
    global _VAD_MODEL
    if _VAD_MODEL is not None:
        return _VAD_MODEL
    with _VAD_LOCK:
        if _VAD_MODEL is None:
            try:
                import torch
                from silero_vad import load_silero_vad
                _VAD_MODEL = (torch, load_silero_vad())
            except Exception:
                _VAD_MODEL = False
    return _VAD_MODEL


def _vad_energy(y, sr: int, energy_db_thresh: float, min_silence: float,
                min_speech: float):
    """Energy-based VAD (fallback when Silero is unavailable)."""
    frame_len = int(sr * FRAME_MS / 1000)
    hop = int(sr * HOP_MS / 1000)
    rms = librosa.feature.rms(y=y, frame_length=frame_len, hop_length=hop)[0]
    voiced = 20 * np.log10(rms + 1e-10) > energy_db_thresh

    raw = []
    start = None
    for i, is_voice in enumerate(voiced):
        t = i * hop / sr
        if is_voice and start is None:
            start = t
        elif not is_voice and start is not None:
            if t - start >= min_speech:
                raw.append((start, t))
            start = None
    if start is not None and len(y) / sr - start >= min_speech:
        raw.append((start, len(y) / sr))

    # Merge segments separated by a short silence (same utterance / interjection).
    merged = []
    for seg in raw:
        if merged and seg[0] - merged[-1][1] <= min_silence:
            merged[-1] = (merged[-1][0], seg[1])
        else:
            merged.append(seg)
    return merged


def vad_segments(y, sr: int, energy_db_thresh: float = ENERGY_DB_THRESH,
                 min_silence: float = MIN_SILENCE_S,
                 min_speech: float = MIN_SPEECH_S,
                 vad: str = "auto"):
    """Return speech segments as ``[(start_s, end_s), ...]``.

    ``vad`` selects the backend: "auto" uses Silero when installed and falls
    back to the energy-based detector, "silero" forces Silero, "energy"
    forces the lightweight energy-based detector.
    """
    if vad != "energy":
        silero = _get_silero_vad()
        if silero is not False:
            torch, model = silero
            y16 = librosa.resample(y, orig_sr=sr, target_sr=VAD_SR)
            audio = torch.from_numpy(np.ascontiguousarray(y16, dtype=np.float32)).float()
            try:
                from silero_vad import get_speech_timestamps
                timestamps = get_speech_timestamps(
                    audio, model, sampling_rate=VAD_SR, return_seconds=True,
                    min_speech_duration_ms=int(min_speech * 1000),
                    min_silence_duration_ms=int(min_silence * 1000),
                    speech_pad_ms=200,  # context padding: don't chop word edges
                )
                segments = [(float(t["start"]), float(t["end"])) for t in timestamps]
                if segments:
                    return segments
            except Exception:
                pass  # fall through to the energy-based VAD
    return _vad_energy(y, sr, energy_db_thresh, min_silence, min_speech)


def _segment_features(y, sr: int, segments):
    """Mean MFCC vector per speech segment; returns (X, valid_mask)."""
    xs = []
    valid = []
    for seg in segments:
        i0, i1 = int(seg[0] * sr), int(seg[1] * sr)
        if i1 - i0 < int(MIN_SEGMENT_S * sr):
            xs.append(None)
            valid.append(False)
            continue
        mfcc = librosa.feature.mfcc(y=y[i0:i1], sr=sr, n_mfcc=13).mean(axis=1)
        xs.append(mfcc)
        valid.append(True)
    X = np.array([x for x in xs if x is not None])
    return X, np.array(valid)


def _speechbrain_available() -> bool:
    """Lazily import speechbrain/torch; return whether ECAPA embeddings are usable."""
    global _HAS_SPEECHBRAIN, _TORCH, _ENCODER_CLASSIFIER
    if _HAS_SPEECHBRAIN is None:
        try:
            import torch
            try:
                #from speechbrain.inference.speaker import EncoderClassifier
                from speechbrain.inference import EncoderClassifier
            except ImportError:  # older speechbrain (< 1.0)
                from speechbrain.pretrained import EncoderClassifier
            _TORCH, _ENCODER_CLASSIFIER = torch, EncoderClassifier
            _HAS_SPEECHBRAIN = True
        except Exception:
            _HAS_SPEECHBRAIN = False
    return _HAS_SPEECHBRAIN


_ecapa_model = None
_ecapa_lock = threading.Lock()


def _get_ecapa_model():
    """Load (and cache) the ECAPA-TDNN speaker encoder; downloads on first use."""
    global _ecapa_model
    if not _speechbrain_available():
        raise RuntimeError("speechbrain is not installed")
    if _ecapa_model is not None:
        return _ecapa_model
    with _ecapa_lock:
        if _ecapa_model is None:
            _ecapa_model = _ENCODER_CLASSIFIER.from_hparams(
                source=ECAPA_SOURCE, savedir=ECAPA_SAVEDIR
            )
    return _ecapa_model


def _embed_segments(y16, segments, model, min_seconds: float = ASSIGN_MIN_S):
    """ECAPA embedding per speech segment (``None`` when too short to embed).

    Segments shorter than ``EMBEDDING_MIN_S`` are still embedded so they can be
    *assigned* to a speaker afterwards, even though they don't drive clustering.
    """
    embs = []
    for seg in segments:
        i0, i1 = int(seg[0] * EMBEDDING_SR), int(seg[1] * EMBEDDING_SR)
        if i1 - i0 < int(min_seconds * EMBEDDING_SR):
            embs.append(None)
            continue
        wav = _TORCH.from_numpy(y16[i0:i1]).float().unsqueeze(0)
        with _TORCH.no_grad():
            emb = model.encode_batch(wav, wav_lens=_TORCH.ones(1))
        embs.append(emb.squeeze(0).mean(dim=0).numpy())
    return embs


def _l2_normalize(E):
    """Row-wise L2 normalisation (ECAPA embeddings are scored by cosine)."""
    return E / (np.linalg.norm(E, axis=1, keepdims=True) + 1e-8)


def _auto_cluster_labels(E, max_speakers: int = MAX_SPEAKERS,
                         min_per_cluster: int = 2):
    """Pick the speaker count by silhouette score over KMeans.

    Embeddings are L2-normalised first so the Euclidean distance KMeans
    minimises matches the cosine metric the speaker model was trained with.
    Tries every k in [2, ``max_speakers``] and keeps the split with the best
    silhouette, skipping any k where a cluster would be too small. Returns all
    zeros (a single actor) when no split clearly wins.
    """
    E = _l2_normalize(E)
    n = E.shape[0]
    if n < min_per_cluster * 2:
        return np.zeros(n, dtype=int)
    best_k, best_score = 1, -1.0
    for k in range(2, min(max_speakers, n // min_per_cluster) + 1):
        km = KMeans(n_clusters=k, n_init=10, random_state=42)
        labels = km.fit_predict(E)
        counts = np.bincount(labels, minlength=k)
        if counts.min() < min_per_cluster:  # tiny clusters = over-split
            continue
        score = float(silhouette_score(E, labels, metric="cosine"))
        if score > best_score:
            best_k, best_score = k, score
    if best_k == 1:
        return np.zeros(n, dtype=int)
    return KMeans(n_clusters=best_k, n_init=10, random_state=42).fit_predict(E)


def _cluster_mfcc(y, sr: int, segments, n_speakers=None):
    """Cluster speech segments from MFCCs (fallback when embeddings are unavailable)."""
    X, valid = _segment_features(y, sr, segments)
    labels = np.zeros(len(segments), dtype=int)
    if not _HAS_SKLEARN or X.shape[0] == 0:
        return labels
    # Standardise each MFCC dimension so coefficient 0 doesn't dominate.
    mean, std = X.mean(axis=0), X.std(axis=0) + 1e-8
    Xs = (X - mean) / std
    if n_speakers is not None:
        k = max(1, min(n_speakers, X.shape[0]))
        if k > 1:
            labels[valid] = KMeans(n_clusters=k, n_init=10, random_state=42).fit_predict(Xs)
        return labels
    labels[valid] = _auto_cluster_labels(Xs)
    return labels


def _merge_windows(wins, gap: float = 0.35):
    """Merge consecutive same-speaker windows into contiguous segments."""
    merged = []
    for w in sorted(wins):
        if merged and w[0] <= merged[-1][1] + gap:
            merged[-1] = (merged[-1][0], max(merged[-1][1], w[1]))
        else:
            merged.append(w)
    return merged


def _window_segments(y, sr: int, segments, model, n_speakers=None):
    """Window-based ECAPA classification.

    Reliable segments (>= ``EMBEDDING_MIN_S``) are clustered into speaker
    centroids, then every VAD segment is scanned with ``WINDOW_S``-long windows
    (``WINDOW_HOP_S`` stride). Each window is assigned to its nearest centroid;
    a window whose top-2 centroid similarity gap is below ``OVERLAP_GAP`` is
    assigned to *both* speakers (overlapped speech). The primary label is
    median-filtered over the neighbouring windows to remove single-window
    flicker, and overlap tiles are split between the two actors by their
    similarity share so the tracks keep summing back to the original audio.

    Returns ``{speaker: (merged_segments, weighted_tiles)}`` where
    ``weighted_tiles`` is ``[(start, end, weight), ...]`` (the units used to
    build the isolated tracks) and ``merged_segments`` is the display version.
    """
    y16 = librosa.resample(y, orig_sr=sr, target_sr=EMBEDDING_SR)
    embs = _embed_segments(y16, segments, model)
    rel = [i for i, e in enumerate(embs)
           if e is not None and segments[i][1] - segments[i][0] >= EMBEDDING_MIN_S]
    if not rel:
        return {0: (list(segments), _as_tiles(segments))}

    E = _l2_normalize(np.stack([embs[i] for i in rel]))
    if n_speakers is not None:
        k = max(1, min(n_speakers, len(rel)))
        if k == 1:
            return {0: (list(segments), _as_tiles(segments))}
        rel_labels = KMeans(n_clusters=k, n_init=10, random_state=42).fit_predict(E)
    else:
        rel_labels = _auto_cluster_labels(E)

    n_clusters = int(rel_labels.max()) + 1
    centroids = _l2_normalize(
        np.stack([E[rel_labels == c].mean(axis=0) for c in range(n_clusters)])
    )

    # --- window scan (each window owns the leading WINDOW_HOP_S tile) ---
    wins = []  # (start, tile_end, primary, secondary_or_None, sims)
    for s, e in segments:
        t = s
        while t < e - 0.05:
            w_len = min(WINDOW_S, e - t)
            if w_len < ASSIGN_MIN_S:
                break
            i0, i1 = int(t * EMBEDDING_SR), int((t + w_len) * EMBEDDING_SR)
            wav = _TORCH.from_numpy(y16[i0:i1]).float().unsqueeze(0)
            with _TORCH.no_grad():
                emb = model.encode_batch(wav, wav_lens=_TORCH.ones(1)).squeeze(0).mean(dim=0).numpy()
            emb = emb / (np.linalg.norm(emb) + 1e-8)
            sims = centroids @ emb
            order = np.argsort(sims)[::-1]
            primary, second = int(order[0]), None
            if n_speakers is None and n_clusters >= 2 \
                    and sims[order[0]] - sims[order[1]] < OVERLAP_GAP:
                second = int(order[1])
            tile_end = min(t + WINDOW_HOP_S, e)
            wins.append((t, tile_end, primary, second, sims))
            t += WINDOW_HOP_S

    # --- median filter on the primary label (kills single-window flicker) ---
    prim = [w[2] for w in wins]
    smoothed = list(prim)
    for i in range(len(prim)):
        lo, hi = max(0, i - 1), min(len(prim), i + 2)
        counts = {}
        for j in range(lo, hi):
            counts[prim[j]] = counts.get(prim[j], 0) + 1
        best = max(counts, key=counts.get)
        if counts[best] >= 2:
            smoothed[i] = best

    # --- per-speaker tiles with similarity-weighted overlap splits ---
    per_windows = {c: [] for c in range(n_clusters)}
    for i, (t, tile_end, _, second, sims) in enumerate(wins):
        p = smoothed[i]
        q = second if (second is not None and second != p) else None
        if q is None:
            per_windows[p].append((t, tile_end, 1.0))
        else:
            sp_, sq_ = float(sims[p]), float(sims[q])
            wq = 1.0 / (1.0 + np.exp(sp_ - sq_))  # softmax share over the top-2
            per_windows[p].append((t, tile_end, 1.0 - wq))
            per_windows[q].append((t, tile_end, wq))

    per_speaker = {}
    for c in range(n_clusters):
        wins_c = per_windows[c]
        merged = _merge_windows([(t, e) for t, e, _ in wins_c])
        if merged:
            per_speaker[c] = (merged, wins_c)
    return per_speaker


# --------------------------------------------------------------------------- #
# Optional pyannote-audio 3.1 diarization engine (state-of-the-art).
# Needs a Hugging Face token with access to the gated model.
# --------------------------------------------------------------------------- #
_PYANNOTE = None  # (token, pipeline_or_None, timestamp)
_PYANNOTE_LOCK = threading.Lock()
_PYANNOTE_FAIL_TTL = 60.0  # retry a rejected token after this long (license may have been accepted)
PYANNOTE_MODEL = "pyannote/speaker-diarization-3.1"
PYANNOTE_SR = 16000

# Which engine the last ``diarize`` call actually used (for the UI to report):
#   "pyannote" | "ecapa" | "mfcc" | None
last_engine_used = None


def _get_pyannote_pipeline(token: str):
    """Load (and cache) the pyannote 3.1 pipeline for ``token``.

    A rejected token is only remembered for ``_PYANNOTE_FAIL_TTL`` seconds, so a
    license accepted mid-session is picked up without restarting the app.
    """
    global _PYANNOTE
    now = time.monotonic()
    if _PYANNOTE is not None and _PYANNOTE[0] == token:
        pipe, ts = _PYANNOTE[1], _PYANNOTE[2]
        if pipe is not None or now - ts < _PYANNOTE_FAIL_TTL:
            return pipe
    with _PYANNOTE_LOCK:
        if _PYANNOTE is not None and _PYANNOTE[0] == token:
            pipe, ts = _PYANNOTE[1], _PYANNOTE[2]
            if pipe is not None or now - ts < _PYANNOTE_FAIL_TTL:
                return pipe
        try:
            from pyannote.audio import Pipeline
            pipeline = Pipeline.from_pretrained(PYANNOTE_MODEL, token=token)
        except Exception:
            _PYANNOTE = (token, None, time.monotonic())  # remember the failure
            return None
        _PYANNOTE = (token, pipeline, time.monotonic())
        return pipeline


def _diarize_pyannote(y, sr: int, n_speakers=None, token: str | None = None):
    """Run the pyannote 3.1 pipeline; return ``None`` on any failure."""
    if not token:
        return None
    try:
        pipeline = _get_pyannote_pipeline(token)
        if pipeline is None:
            return None
        import torch
        y16 = librosa.resample(y, orig_sr=sr, target_sr=PYANNOTE_SR)
        waveform = torch.from_numpy(
            np.ascontiguousarray(y16, dtype=np.float32)
        ).float().unsqueeze(0)
        kwargs = {"num_speakers": n_speakers} if n_speakers is not None else {}
        annotation = pipeline({"waveform": waveform, "sample_rate": PYANNOTE_SR}, **kwargs)
        # pyannote >= 4 wraps the Annotation in a DiarizeOutput — unwrap it.
        if hasattr(annotation, "speaker_diarization"):
            annotation = annotation.speaker_diarization

        per = {}
        for turn, _, label in annotation.itertracks(yield_label=True):
            if turn.end - turn.start >= MIN_SEGMENT_S:  # drop sub-0.5s artifacts
                per.setdefault(label, []).append((turn.start, turn.end))
        if not per:
            return None

        per_speaker = {}
        # Actor 1 = the most talkative voice.
        for i, label in enumerate(sorted(per, key=lambda l: -sum(e - s for s, e in per[l]))):
            segs = _merge_windows(sorted(per[label]), gap=0.25)
            per_speaker[i] = (segs, _as_tiles(segs))
        return per_speaker
    except Exception:
        return None  # caller falls back to ECAPA / MFCC


def _as_tiles(segments, hop: float = WINDOW_HOP_S):
    """Split segments into hop-length tiles ``(start, end, 1.0)``.

    Uniform building blocks for the per-actor tracks: consecutive tiles are
    disjoint, so track construction can accumulate without double-counting.
    """
    tiles = []
    for s, e in segments:
        t = s
        while t < e:
            tiles.append((t, min(t + hop, e), 1.0))
            t += hop
    return tiles


def diarize(y, sr: int, n_speakers: int | None = None, engine: str = "auto",
            hf_token: str | None = None, vad: str = "auto"):
    """Split speech into one actor slot per detected voice.

    ``engine`` selects the diarization backend:
      - "auto": pyannote 3.1 when a usable ``hf_token`` is provided, else
        ECAPA-TDNN window classification (speechbrain), else MFCC clustering.
      - "pyannote": force pyannote 3.1 (falls back to ECAPA / MFCC on failure).
      - "ecapa": ECAPA-TDNN window classification (MFCC fallback).

    ``vad`` is passed through to :func:`vad_segments` ("auto" / "silero" /
    "energy"). Sets the module-level ``last_engine_used`` so the UI can report
    which engine ran (and warn about degraded quality). Returns a list of
    speaker dicts:
        {id, label, segments: [(start, end)...],
         windows: [(start, end, weight)...], track: np.ndarray}
    """
    global last_engine_used
    last_engine_used = None

    per_speaker = None
    if engine in ("auto", "pyannote"):
        per_speaker = _diarize_pyannote(y, sr, n_speakers, hf_token)
        if per_speaker is not None:
            last_engine_used = "pyannote"

    if per_speaker is None:
        segments = vad_segments(y, sr, vad=vad)
        if segments and _speechbrain_available():
            try:
                model = _get_ecapa_model()
                per_speaker = _window_segments(y, sr, segments, model, n_speakers)
                last_engine_used = "ecapa"
            except Exception:
                # Model download/load failed (e.g. offline) — MFCC fallback.
                per_speaker = None
        if per_speaker is None and segments:
            labels = _cluster_mfcc(y, sr, segments, n_speakers)
            per_speaker = {
                k: ([segments[i] for i in np.where(labels == k)[0]],
                    _as_tiles([segments[i] for i in np.where(labels == k)[0]]))
                for k in range(int(labels.max()) + 1)
            }
            last_engine_used = "mfcc"

    if not per_speaker:
        return []

    speakers = []
    for k in sorted(per_speaker):
        segs, windows = per_speaker[k]
        segs = sorted(segs)
        if not segs:
            continue
        speakers.append({
            "id": f"actor_{k + 1}",
            "label": f"Actor {k + 1}",
            "segments": segs,
            "windows": windows,
        })
    # Actor 1 = the most talkative voice (stable ordering across runs).
    speakers.sort(key=lambda s: -sum(e - st for st, e in s["segments"]))
    for i, sp in enumerate(speakers):
        sp["id"] = f"actor_{i + 1}"
        sp["label"] = f"Actor {i + 1}"
    _build_tracks(speakers, y, sr)
    return speakers


def _build_tracks(speakers, y, sr: int) -> None:
    """Give each speaker an isolated track from their weighted tiles.

    Overlap tiles are split between the two actors by their similarity share
    (weights sum to 1), so the per-actor tracks sum back to the original audio
    and the mix never double-applies gain on overlapping speech.
    """
    for sp in speakers:
        track = np.zeros_like(y)
        for start, end, weight in sp["windows"]:
            i0, i1 = int(start * sr), int(end * sr)
            track[i0:i1] += y[i0:i1] * weight
        sp["track"] = track


# --------------------------------------------------------------------------- #
# 3. Speaker profiling (gender & age heuristics based on pitch)
# --------------------------------------------------------------------------- #
def _median_f0(speaker, sr: int, max_seconds: float = 30.0):
    """Median fundamental frequency over (up to ``max_seconds`` of) the actor's speech."""
    parts = []
    total = 0.0
    for start, end in speaker["segments"]:
        if total >= max_seconds:
            break
        i0, i1 = int(start * sr), int(end * sr)
        parts.append(speaker["track"][i0:i1])
        total += end - start
    if not parts:
        return None
    speech = np.concatenate(parts)
    f0 = librosa.yin(speech, fmin=75.0, fmax=600.0, sr=sr)
    f0 = f0[~np.isnan(f0)]
    if f0.size == 0:
        return None
    return float(np.median(f0))


def _guess_gender(f0: float) -> str:
    if f0 < 160.0:
        return "Male"
    if f0 > 185.0:
        return "Female"
    return "Unknown"


def _guess_age(f0: float) -> str:
    if f0 > 240.0:
        return "Child"
    if f0 < 115.0:
        return "Elderly"
    return "Adult"


def _get_age_gender_model():
    """Lazily load the audeering wav2vec2 age/gender model (audonnx + ONNX).

    Downloads the ONNX model once into a temp dir on first use. Returns ``False``
    when audonnx isn't installed or the download fails, so callers fall back to
    the pitch heuristics.
    """
    global _AGE_GENDER_MODEL
    if _AGE_GENDER_MODEL is not None:
        return _AGE_GENDER_MODEL
    with _AGE_GENDER_LOCK:
        if _AGE_GENDER_MODEL is not None:
            return _AGE_GENDER_MODEL
        try:
            import audonnx  # noqa: F401
        except Exception:
            _AGE_GENDER_MODEL = False
            return False
        try:
            model_root = os.path.join(AGE_GENDER_SAVEDIR, "model")
            os.makedirs(model_root, exist_ok=True)
            if not any(f.endswith(".onnx") for f in os.listdir(model_root)):
                zip_path = os.path.join(AGE_GENDER_SAVEDIR, "age_gender.zip")
                if not os.path.exists(zip_path):
                    import urllib.request
                    urllib.request.urlretrieve(AGE_GENDER_URL, zip_path)
                import zipfile
                with zipfile.ZipFile(zip_path) as z:
                    z.extractall(model_root)
            _AGE_GENDER_MODEL = audonnx.load(model_root)
        except Exception:
            _AGE_GENDER_MODEL = False
    return _AGE_GENDER_MODEL


def _speech_16k(speaker, sr: int, max_seconds: float = 30.0):
    """Up to ``max_seconds`` of the actor's speech, resampled to 16 kHz."""
    parts, total = [], 0.0
    for start, end in speaker["segments"]:
        if total >= max_seconds:
            break
        i0, i1 = int(start * sr), int(end * sr)
        parts.append(speaker["track"][i0:i1])
        total += end - start
    if not parts:
        return None
    speech = np.concatenate(parts)
    return librosa.resample(speech, orig_sr=sr, target_sr=AGE_GENDER_SR)


def _predict_age_gender(speech16, model):
    """Run the audeering model -> ``(age_years, gender_label)``."""
    out = model(np.asarray(speech16, dtype=np.float32), AGE_GENDER_SR)
    age = float(np.asarray(out["logits_age"]).reshape(-1)[0]) * 100.0
    gender = AGE_GENDER_LABELS[int(np.argmax(np.asarray(out["logits_gender"]).reshape(-1)))]
    return age, gender


def _age_bucket(age_years: float) -> str:
    if age_years < 14:
        return "Child"
    if age_years >= 60:
        return "Elderly"
    return "Adult"


def profile_speakers(speakers, sr: int, estimator: str = "auto") -> None:
    """Fill each speaker dict with loudness, speech stats and gender/age estimates.

    Gender & age come from the audeering wav2vec2 model (age in years + one of
    female/male/child) when audonnx is available; otherwise they fall back to
    the lightweight pitch heuristics. ``estimator`` selects the backend:
    "auto" (audeering when available, else heuristics), "audeering" (force the
    ONNX model, heuristics on failure), or "heuristics" (instant, offline).
    """
    model = _get_age_gender_model() if estimator != "heuristics" else False
    for sp in speakers:
        track = sp["track"]
        sp["rms_db"] = _rms_db(track)
        sp["speech_seconds"] = float(sum(e - s for s, e in sp["segments"]))
        f0 = _median_f0(sp, sr)
        sp["f0_median"] = f0

        if model is not False:
            speech16 = _speech_16k(sp, sr)
            if speech16 is not None:
                try:
                    age, gender = _predict_age_gender(speech16, model)
                    sp["age_years"] = round(age, 1)
                    sp["gender_auto"] = "Unknown" if gender == "Child" else gender
                    sp["age_auto"] = _age_bucket(age)
                except Exception:  # inference hiccup — fall back to heuristics
                    sp["gender_auto"] = _guess_gender(f0) if f0 is not None else "Unknown"
                    sp["age_auto"] = _guess_age(f0) if f0 is not None else "Unknown"
            else:
                sp["gender_auto"] = _guess_gender(f0) if f0 is not None else "Unknown"
                sp["age_auto"] = _guess_age(f0) if f0 is not None else "Unknown"
        else:
            sp["gender_auto"] = _guess_gender(f0) if f0 is not None else "Unknown"
            sp["age_auto"] = _guess_age(f0) if f0 is not None else "Unknown"
        sp["preview_bytes"] = wav_bytes(track, sr)


def _rms_db(track) -> float:
    rms = float(np.sqrt(np.mean(track ** 2)))
    return 20.0 * np.log10(rms + 1e-12)


# --------------------------------------------------------------------------- #
# 4. Gain adjustment & mixing
# --------------------------------------------------------------------------- #
def auto_gains(speakers, target_db: float = -20.0, max_gain: float = 20.0):
    """Equal-loudness gains (dB) so every actor sits at ``target_db`` RMS."""
    return {
        sp["id"]: float(np.clip(target_db - sp["rms_db"], -max_gain, max_gain))
        for sp in speakers
    }


def mix(speakers, gains, original, sr: int):
    """Mix actor tracks with their dB gains over the residual (non-speech) audio."""
    base = np.zeros_like(original)
    for sp in speakers:
        base = base + sp["track"]
    residual = original - base  # background / music / noise not attributed to any actor

    out = np.zeros_like(original)
    for sp in speakers:
        g = 10.0 ** (float(gains.get(sp["id"], 0.0)) / 20.0)
        out = out + sp["track"] * g
    out = out + residual

    peak = float(np.max(np.abs(out))) if out.size else 0.0
    if peak > 0.99:
        out = out * (0.99 / peak)
    return out


def wav_bytes(audio, sr: int, subtype: str = "PCM_16") -> bytes:
    """Encode a float array as WAV bytes (for Streamlit playback / download)."""
    buf = io.BytesIO()
    sf.write(buf, audio, sr, format="WAV", subtype=subtype)
    return buf.getvalue()