pipeline_runner.py

#!/usr/bin/env python3
"""Timed extraction pipeline used by the FastAPI app and benchmarks."""

from __future__ import annotations

import hashlib
import importlib.util
import json
import os
import shutil
import subprocess
import sys
import tempfile
import time
from contextlib import contextmanager
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Callable

import librosa
import numpy as np
import soundfile as sf

from sample_extractor import (
    DEMUCS_MODELS,
    DEMUCS_STEMS,
    build_archive,
    classify_hits,
    cluster_hits,
    cluster_hits_online,
    detect_bpm,
    detect_onsets,
    export_midi,
    extract_stem,
    render_midi_with_samples,
    sample_quality_score,
    select_best,
    synthesize_from_cluster,
)

ProgressCallback = Callable[[dict[str, Any]], None]

SPLEETER_MODELS = ["spleeter:4stems", "spleeter:2stems", "spleeter:5stems"]
SPLEETER_STEMS = {
    "spleeter:2stems": ["vocals", "accompaniment"],
    "spleeter:4stems": ["vocals", "drums", "bass", "other"],
    "spleeter:5stems": ["vocals", "drums", "bass", "piano", "other"],
}
SEPARATION_BACKENDS = ["spleeter", "demucs", "none"]


def _module_available(name: str) -> bool:
    return importlib.util.find_spec(name) is not None


def separation_runtime_status() -> dict[str, Any]:
    """Return lightweight runtime diagnostics for optional separation engines.

    This is intentionally import-spec/CLI based and does not instantiate models or
    download weights. It is safe to call on page load.
    """
    spleeter_available = _module_available("spleeter") or shutil.which("spleeter") is not None
    demucs_available = _module_available("demucs")
    return {
        "backends": {
            "spleeter": {
                "available": bool(spleeter_available),
                "default": True,
                "install_hint": "Install requirements-spleeter.txt or use the full-mix fallback.",
            },
            "demucs": {
                "available": bool(demucs_available),
                "default": False,
                "install_hint": "Install demucs and model dependencies for quality fallback separation.",
            },
            "none": {
                "available": True,
                "default": False,
                "install_hint": "Always available; uses the uploaded mix without source separation.",
            },
        }
    }

@dataclass
class PipelineParams:
    stem: str = "drums"
    separation_backend: str = "spleeter"
    spleeter_model: str = "spleeter:4stems"
    demucs_model: str = "htdemucs_ft"
    demucs_shifts: int = 1
    demucs_overlap: float = 0.25
    onset_mode: str = "auto"
    onset_delta: float = 0.12
    energy_threshold_db: float = -35.0
    pre_pad: float = 0.003
    min_dur: float = 0.02
    max_dur: float = 1.5
    min_gap: float = 0.03
    ncc_threshold: float = 0.80
    attack_ms: float = 25.0
    mel_threshold: float = 0.75
    linkage: str = "average"
    clustering_mode: str = "batch_quality"
    target_min: int = 5
    target_max: int = 20
    synthesize: bool = True
    quantize_midi: bool = True
    subdivision: int = 16
    device: str = "cpu"
    auto_tune: bool = True
    use_disk_cache: bool = True
    allow_backend_fallback: bool = True

    @classmethod
    def from_mapping(cls, data: dict[str, Any] | None) -> "PipelineParams":
        data = dict(data or {})
        allowed = {field.name for field in cls.__dataclass_fields__.values()}
        unknown = sorted(set(data) - allowed)
        if unknown:
            raise ValueError(f"Unknown pipeline parameter(s): {', '.join(unknown)}")

        int_fields = {"demucs_shifts", "target_min", "target_max", "subdivision"}
        float_fields = {
            "demucs_overlap",
            "onset_delta",
            "energy_threshold_db",
            "pre_pad",
            "min_dur",
            "max_dur",
            "min_gap",
            "ncc_threshold",
            "attack_ms",
            "mel_threshold",
        }
        bool_fields = {"synthesize", "quantize_midi", "auto_tune", "use_disk_cache", "allow_backend_fallback"}

        def coerce_bool(name: str, value: Any) -> bool:
            if isinstance(value, bool):
                return value
            if isinstance(value, (int, float)) and value in {0, 1}:
                return bool(value)
            if isinstance(value, str):
                normalized = value.strip().lower()
                if normalized in {"true", "1", "yes", "on"}:
                    return True
                if normalized in {"false", "0", "no", "off"}:
                    return False
            raise ValueError(f"{name} must be a boolean")

        coerced: dict[str, Any] = {}
        for name, value in data.items():
            if value is None:
                coerced[name] = value
            elif name in int_fields:
                try:
                    number = float(value) if isinstance(value, str) else value
                    if int(number) != number:
                        raise ValueError
                    coerced[name] = int(number)
                except (TypeError, ValueError):
                    raise ValueError(f"{name} must be an integer") from None
            elif name in float_fields:
                try:
                    coerced[name] = float(value)
                except (TypeError, ValueError):
                    raise ValueError(f"{name} must be a number") from None
            elif name in bool_fields:
                coerced[name] = coerce_bool(name, value)
            else:
                coerced[name] = value

        params = cls(**coerced)
        params.validate()
        return params

    def validate(self) -> None:
        if self.separation_backend not in set(SEPARATION_BACKENDS):
            raise ValueError(f"Unsupported separation backend: {self.separation_backend}")
        if self.spleeter_model not in SPLEETER_MODELS:
            raise ValueError(f"Unsupported Spleeter model: {self.spleeter_model}")
        if self.demucs_model not in DEMUCS_MODELS:
            raise ValueError(f"Unsupported Demucs model: {self.demucs_model}")
        if self.separation_backend == "demucs":
            allowed_stems = set(DEMUCS_STEMS.get(self.demucs_model, [])) | {"all"}
            backend_label = self.demucs_model
        elif self.separation_backend == "spleeter":
            allowed_stems = set(SPLEETER_STEMS.get(self.spleeter_model, [])) | {"all"}
            backend_label = self.spleeter_model
        else:
            allowed_stems = {"all"}
            backend_label = "full mix"
        if self.stem not in allowed_stems:
            raise ValueError(f"Stem '{self.stem}' is not available for {backend_label}")
        if self.onset_mode not in {"auto", "percussive", "harmonic", "broadband"}:
            raise ValueError(f"Unsupported onset mode: {self.onset_mode}")
        if self.linkage not in {"average", "complete", "single"}:
            raise ValueError(f"Unsupported clustering linkage: {self.linkage}")
        if self.clustering_mode not in {"batch_quality", "online_preview"}:
            raise ValueError(f"Unsupported clustering mode: {self.clustering_mode}")
        if not 0 <= self.demucs_shifts <= 8:
            raise ValueError("demucs_shifts must be between 0 and 8")
        if not 0.0 <= self.demucs_overlap <= 0.9:
            raise ValueError("demucs_overlap must be between 0.0 and 0.9")
        if not 0.001 <= self.onset_delta <= 1.0:
            raise ValueError("onset_delta must be between 0.001 and 1.0")
        if not -100.0 <= self.energy_threshold_db <= 0.0:
            raise ValueError("energy_threshold_db must be between -100 and 0 dB")
        if not 0.0 <= self.pre_pad <= 0.25:
            raise ValueError("pre_pad must be between 0 and 0.25 seconds")
        if not 0.001 <= self.min_dur <= self.max_dur <= 10.0:
            raise ValueError("duration bounds must satisfy 0.001 <= min_dur <= max_dur <= 10")
        if not 0.001 <= self.min_gap <= 1.0:
            raise ValueError("min_gap must be between 0.001 and 1.0 seconds")
        if not 0.0 <= self.ncc_threshold <= 1.0:
            raise ValueError("ncc_threshold must be between 0 and 1")
        if not 1.0 <= self.attack_ms <= 250.0:
            raise ValueError("attack_ms must be between 1 and 250 ms")
        if not 0.0 <= self.mel_threshold <= 1.0:
            raise ValueError("mel_threshold must be between 0 and 1")
        if self.target_min < 0 or self.target_max < 0:
            raise ValueError("target_min and target_max must be non-negative")
        if self.target_max and self.target_min and self.target_min > self.target_max:
            raise ValueError("target_min cannot be greater than target_max")
        if self.subdivision not in {4, 8, 16, 32, 64}:
            raise ValueError("subdivision must be one of 4, 8, 16, 32, 64")


@dataclass
class StageTiming:
    key: str
    label: str
    duration_sec: float = 0.0
    status: str = "pending"
    detail: str = ""
    progress: float = 0.0
    work_done: int | None = None
    work_total: int | None = None


@dataclass
class PipelineResult:
    params: dict[str, Any]
    duration_sec: float
    audio_duration_sec: float
    realtime_factor: float
    bpm: float | None
    sample_rate: int
    hit_count: int
    cluster_count: int
    stages: list[dict[str, Any]]
    samples: list[dict[str, Any]]
    hits: list[dict[str, Any]]
    overview: dict[str, Any]
    files: dict[str, str]


STAGE_DEFS = [
    ("stem", "Stem separation / source load"),
    ("auto_tune", "Automatic parameter tuning"),
    ("bpm", "Tempo detection"),
    ("onsets", "Onset detection + slicing"),
    ("classification", "Spectral rule classification"),
    ("clustering", "Mel fingerprint + transient NCC clustering"),
    ("selection", "Best representative scoring"),
    ("synthesis", "Optional sample synthesis"),
    ("export", "MIDI, reproduced mix, WAV, ZIP export"),
]


def initial_stages() -> list[dict[str, Any]]:
    return [asdict(StageTiming(key=key, label=label)) for key, label in STAGE_DEFS]


def _notify(cb: ProgressCallback | None, payload: dict[str, Any]) -> None:
    if cb:
        cb(payload)


def _progress_payload(stages: list[StageTiming]) -> dict[str, Any]:
    running = next((stage for stage in stages if stage.status == "running"), None)
    running_progress = float(np.clip(running.progress, 0.0, 1.0)) if running else 0.0
    total = 0.0
    completed_units = 0.0
    completed = 0
    for stage in stages:
        units = float(stage.work_total or 1)
        total += units
        if stage.status == "done":
            completed += 1
            completed_units += units
        elif stage.status == "running":
            completed_units += float(stage.work_done) if stage.work_total and stage.work_done is not None else running_progress * units
    fraction = float(np.clip(completed_units / max(total, 1.0), 0.0, 1.0))
    return {
        "fraction": round(fraction, 6),
        "completed_steps": completed,
        "total_steps": len(stages),
        "completed_units": round(completed_units, 6),
        "total_units": round(total, 6),
        "stage_key": running.key if running else None,
        "stage_label": running.label if running else None,
        "stage_fraction": round(running_progress, 6),
        "stage_work_done": running.work_done if running else None,
        "stage_work_total": running.work_total if running else None,
        "basis": "exact completed work units: Demucs chunks when available; Spleeter and non-instrumented stages advance only at real stage boundaries; no time-based estimates",
    }


def _notify_stages(cb: ProgressCallback | None, stages: list[StageTiming], event_type: str, stage: StageTiming) -> None:
    _notify(cb, {
        "type": event_type,
        "stage": asdict(stage),
        "stages": [asdict(item) for item in stages],
        "progress": _progress_payload(stages),
    })


@contextmanager
def _timed_stage(stages: list[StageTiming], key: str, cb: ProgressCallback | None = None):
    stage = next(stage for stage in stages if stage.key == key)
    stage.status = "running"
    stage.progress = 0.0
    stage.work_done = None
    stage.work_total = None
    _notify_stages(cb, stages, "stage", stage)
    started = time.perf_counter()
    try:
        yield stage
    except Exception as exc:
        stage.duration_sec = time.perf_counter() - started
        stage.status = "error"
        stage.detail = str(exc)
        _notify_stages(cb, stages, "stage", stage)
        raise
    else:
        stage.duration_sec = time.perf_counter() - started
        stage.status = "done"
        stage.progress = 1.0
        _notify_stages(cb, stages, "stage", stage)


def _update_stage_progress(
    stages: list[StageTiming],
    key: str,
    cb: ProgressCallback | None,
    *,
    fraction: float,
    detail: str | None = None,
    work_done: int | None = None,
    work_total: int | None = None,
) -> None:
    stage = next(stage for stage in stages if stage.key == key)
    if stage.status != "running":
        return
    stage.progress = float(np.clip(fraction, 0.0, 1.0))
    stage.work_done = work_done
    stage.work_total = work_total
    if detail:
        stage.detail = detail
    _notify_stages(cb, stages, "stage_progress", stage)


def _normalise_audio(audio: np.ndarray) -> np.ndarray:
    audio = audio.astype(np.float32)
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    peak = float(np.max(np.abs(audio))) if audio.size else 0.0
    if peak > 0:
        audio = audio / peak
    return audio.astype(np.float32)


def _mono(audio: np.ndarray) -> np.ndarray:
    audio = np.asarray(audio, dtype=np.float32)
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    return audio.astype(np.float32)


def _pad_or_trim(audio: np.ndarray, length: int) -> np.ndarray:
    audio = _mono(audio)
    if len(audio) == length:
        return audio
    if len(audio) > length:
        return audio[:length]
    return np.pad(audio, (0, max(0, length - len(audio)))).astype(np.float32)


def _load_source_mix(audio_path: str | os.PathLike[str], sr: int) -> np.ndarray:
    audio, _ = librosa.load(audio_path, sr=sr, mono=True)
    return _mono(audio)


def _common_gain(reference: np.ndarray, fallback: np.ndarray) -> float:
    peak = float(np.max(np.abs(reference))) if reference.size else 0.0
    if peak <= 1e-8:
        peak = float(np.max(np.abs(fallback))) if fallback.size else 0.0
    return peak if peak > 1e-8 else 1.0


def _rms(audio: np.ndarray) -> float:
    audio = _mono(audio)
    return float(np.sqrt(np.mean(np.square(audio, dtype=np.float64)))) if audio.size else 0.0


def _match_rms(rendered: np.ndarray, reference: np.ndarray, *, min_gain: float = 0.05, max_gain: float = 8.0) -> np.ndarray:
    rendered_rms = _rms(rendered)
    reference_rms = _rms(reference)
    if rendered_rms <= 1e-10 or reference_rms <= 1e-10:
        return _mono(rendered)
    gain = float(np.clip(reference_rms / rendered_rms, min_gain, max_gain))
    return (_mono(rendered) * gain).astype(np.float32)


def _soft_limit(audio: np.ndarray, ceiling: float = 0.98) -> np.ndarray:
    audio = _mono(audio).astype(np.float32)
    peak = float(np.max(np.abs(audio))) if audio.size else 0.0
    if peak > ceiling > 0:
        audio = audio * (ceiling / peak)
    return audio.astype(np.float32)


def _make_reproduction_mix(target_reconstruction: np.ndarray, context_bed: np.ndarray, length: int) -> np.ndarray:
    target = _pad_or_trim(target_reconstruction, length)
    context = _pad_or_trim(context_bed, length)
    return _soft_limit(context + target)


MODULE_ROOT = Path(__file__).resolve().parent
CACHE_DIR = Path(os.environ["DSE_CACHE_DIR"]) if os.environ.get("DSE_CACHE_DIR") else MODULE_ROOT / ".cache"
STEM_CACHE_DIR = CACHE_DIR / "stems"
CACHE_VERSION = "dse-cache-v4-runtime-fallback"


def _write_audio(path: Path, audio: np.ndarray, sr: int, subtype: str = "PCM_24") -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    sf.write(path, audio.astype(np.float32), sr, subtype=subtype)


def _sha256_file(path: str | os.PathLike[str]) -> str:
    h = hashlib.sha256()
    with Path(path).open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def _stem_cache_path(audio_path: str | os.PathLike[str], params: PipelineParams) -> Path:
    key_payload = {
        "version": CACHE_VERSION,
        "source_sha256": _sha256_file(audio_path),
        "stem": params.stem,
        "separation_backend": params.separation_backend,
        "spleeter_model": params.spleeter_model,
        "demucs_model": params.demucs_model,
        "demucs_shifts": params.demucs_shifts,
        "demucs_overlap": params.demucs_overlap,
        "device": params.device if params.stem != "all" else "decode",
    }
    key = hashlib.sha256(json.dumps(key_payload, sort_keys=True).encode("utf-8")).hexdigest()
    return STEM_CACHE_DIR / f"{key}.wav"


def _context_cache_path(stem_cache_path: Path) -> Path:
    return stem_cache_path.with_name(f"{stem_cache_path.stem}.context.wav")


def clear_disk_cache() -> None:
    if CACHE_DIR.exists():
        shutil.rmtree(CACHE_DIR)


def _load_spleeter_audio(path: Path, sr: int | None = None) -> tuple[np.ndarray, int]:
    audio, loaded_sr = librosa.load(path, sr=sr, mono=True)
    return _mono(audio), int(loaded_sr)


def _spleeter_output_file(root: Path, source_stem: str, stem: str) -> Path | None:
    candidates = [root / source_stem / f"{stem}.wav", root / f"{stem}.wav"]
    candidates.extend(root.glob(f"**/{stem}.wav"))
    for candidate in candidates:
        if candidate.exists():
            return candidate
    return None


def _extract_spleeter_separation(
    audio_path: str | os.PathLike[str],
    params: PipelineParams,
    progress_cb: Callable[[dict[str, Any]], None] | None = None,
) -> tuple[np.ndarray, int, np.ndarray | None, str]:
    """Run Spleeter and return target stem plus the sum of non-target stems.

    Progress is deliberately coarse because Spleeter does not expose reliable
    chunk callbacks through the CLI/Python API. We report start/completion only;
    the global UI therefore never interpolates fake progress.
    """
    if progress_cb:
        progress_cb({"fraction": 0.0, "completed_units": 0, "total_units": 1, "detail": f"Spleeter {params.spleeter_model} starting"})

    with tempfile.TemporaryDirectory(prefix="dse_spleeter_") as tmp:
        tmpdir = Path(tmp)
        commands = [
            [
                sys.executable,
                "-m",
                "spleeter",
                "separate",
                "-p",
                params.spleeter_model,
                "-o",
                str(tmpdir),
                str(audio_path),
            ],
            [
                "spleeter",
                "separate",
                "-p",
                params.spleeter_model,
                "-o",
                str(tmpdir),
                str(audio_path),
            ],
        ]
        failures: list[str] = []
        completed = None
        for cmd in commands:
            try:
                completed = subprocess.run(cmd, capture_output=True, text=True, check=False, timeout=60 * 30)
            except FileNotFoundError as exc:
                failures.append(str(exc))
                continue
            if completed.returncode == 0:
                break
            failures.append((completed.stderr or completed.stdout or "Spleeter failed").strip()[-1200:])
        if completed is None or completed.returncode != 0:
            raise RuntimeError("; ".join(part for part in failures if part) or "Spleeter failed")

        source_stem = Path(audio_path).stem
        stems = SPLEETER_STEMS[params.spleeter_model]
        paths = {stem: _spleeter_output_file(tmpdir, source_stem, stem) for stem in stems}
        missing = [stem for stem, path in paths.items() if path is None]
        if missing:
            raise RuntimeError(f"Spleeter finished but did not write expected stem(s): {', '.join(missing)}")
        if params.stem not in paths:
            raise RuntimeError(f"Spleeter model {params.spleeter_model} does not provide stem '{params.stem}'")

        target, sr = _load_spleeter_audio(paths[params.stem])  # type: ignore[arg-type]
        context_parts: list[np.ndarray] = []
        for name, path in paths.items():
            if name == params.stem or path is None:
                continue
            part, _ = _load_spleeter_audio(path, sr=sr)
            context_parts.append(_pad_or_trim(part, len(target)))
        context = np.sum(np.stack(context_parts), axis=0).astype(np.float32) if context_parts else None

    if progress_cb:
        progress_cb({"fraction": 1.0, "completed_units": 1, "total_units": 1, "detail": f"Spleeter {params.spleeter_model} complete"})
    return target.astype(np.float32), sr, context, f"{params.stem} via {params.spleeter_model}"


def _extract_demucs_separation(
    audio_path: str | os.PathLike[str],
    params: PipelineParams,
    progress_cb: Callable[[dict[str, Any]], None] | None = None,
) -> tuple[np.ndarray, int, np.ndarray | None, str]:
    audio, sr = extract_stem(
        str(audio_path),
        stem=params.stem,
        device=params.device,
        model_name=params.demucs_model,
        shifts=int(params.demucs_shifts),
        overlap=float(params.demucs_overlap),
        progress_cb=progress_cb,
    )
    return audio, sr, None, f"{params.stem} via Demucs {params.demucs_model}"



def _load_full_mix_separation(
    audio_path: str | os.PathLike[str],
    progress_cb: Callable[[dict[str, Any]], None] | None = None,
    detail: str = "loaded full mix",
) -> tuple[np.ndarray, int, str, np.ndarray | None]:
    audio, sr = librosa.load(audio_path, sr=44100, mono=True)
    if progress_cb:
        progress_cb({"fraction": 1.0, "completed_units": 1, "total_units": 1, "detail": detail})
    return _mono(audio), int(sr), detail, None


def _fallback_to_full_mix(
    audio_path: str | os.PathLike[str],
    progress_cb: Callable[[dict[str, Any]], None] | None,
    reason: str,
) -> tuple[np.ndarray, int, str, np.ndarray | None]:
    return _load_full_mix_separation(
        audio_path,
        progress_cb=progress_cb,
        detail=f"fallback full mix because separation failed: {reason}",
    )

def _load_or_extract_separation(audio_path: str | os.PathLike[str], params: PipelineParams, progress_cb: Callable[[dict[str, Any]], None] | None = None) -> tuple[np.ndarray, int, str, np.ndarray | None]:
    if params.stem == "all" or params.separation_backend == "none":
        return _load_full_mix_separation(audio_path, progress_cb=progress_cb, detail="loaded full mix")

    cache_path = _stem_cache_path(audio_path, params)
    context_cache = _context_cache_path(cache_path)
    if params.use_disk_cache and cache_path.exists():
        audio, sr = sf.read(cache_path, dtype="float32", always_2d=False)
        context = None
        if context_cache.exists():
            context_audio, context_sr = sf.read(context_cache, dtype="float32", always_2d=False)
            context = _pad_or_trim(_mono(context_audio), len(_mono(audio))) if int(context_sr) == int(sr) else _mono(librosa.resample(_mono(context_audio), orig_sr=int(context_sr), target_sr=int(sr)))
        if progress_cb:
            progress_cb({"fraction": 1.0, "completed_units": 1, "total_units": 1, "detail": f"{params.stem} disk-cache hit"})
        return np.asarray(audio, dtype=np.float32), int(sr), f"{params.stem} disk-cache hit", context

    context: np.ndarray | None = None
    if params.separation_backend == "spleeter":
        try:
            audio, sr, context, detail = _extract_spleeter_separation(audio_path, params, progress_cb=progress_cb)
        except Exception as spleeter_exc:
            if not params.allow_backend_fallback:
                raise
            if progress_cb:
                progress_cb({"fraction": 0.0, "completed_units": 0, "total_units": 1, "detail": f"Spleeter unavailable; using full-mix fallback: {spleeter_exc}"})
            audio, sr, detail, context = _fallback_to_full_mix(audio_path, progress_cb, f"Spleeter failed ({spleeter_exc})")
    elif params.separation_backend == "demucs":
        try:
            audio, sr, context, detail = _extract_demucs_separation(audio_path, params, progress_cb=progress_cb)
        except Exception as demucs_exc:
            if not params.allow_backend_fallback:
                raise
            audio, sr, detail, context = _fallback_to_full_mix(audio_path, progress_cb, f"Demucs failed ({demucs_exc})")
    else:
        raise ValueError(f"Unsupported separation backend: {params.separation_backend}")

    if params.use_disk_cache:
        _write_audio(cache_path, audio, sr, subtype="PCM_16")
        if context is not None:
            _write_audio(context_cache, context, sr, subtype="PCM_16")
        detail += " · cached"
    return audio, sr, detail, context



def _build_onset_envelope(audio: np.ndarray, sr: int) -> np.ndarray:
    """Build one onset envelope for automatic tuning; candidate deltas reuse it."""
    y = _mono(audio)
    if y.size == 0:
        return np.zeros(0, dtype=np.float32)
    hop_length = 220
    n_fft = 1024
    try:
        yh, _ = librosa.effects.hpss(y)
        def _sf(part: np.ndarray, fmin: float | None = None, fmax: float | None = None, ms: int = 3, lag: int = 1) -> np.ndarray:
            S = librosa.feature.melspectrogram(
                y=part,
                sr=sr,
                n_fft=n_fft,
                hop_length=hop_length,
                n_mels=96,
                fmin=fmin or 20,
                fmax=fmax or sr // 2,
                power=1,
            )
            return librosa.onset.onset_strength(
                S=librosa.power_to_db(S, ref=np.max),
                sr=sr,
                hop_length=hop_length,
                lag=lag,
                max_size=ms,
            )
        def _norm(values: np.ndarray) -> np.ndarray:
            peak = float(np.max(values)) if values.size else 0.0
            return values / peak if peak > 0 else values
        return np.maximum.reduce([
            _norm(_sf(y, 20, 300)),
            _norm(_sf(y, 300, 4000)),
            _norm(_sf(y, 4000, 16000)),
            _norm(_sf(yh, lag=2, ms=5)),
        ]).astype(np.float32)
    except Exception:
        return np.zeros(0, dtype=np.float32)


def _count_onsets_from_envelope(envelope: np.ndarray, sr: int, delta: float) -> int:
    if envelope.size == 0:
        return 0
    hop_length = 220
    frames = librosa.onset.onset_detect(
        onset_envelope=envelope,
        sr=sr,
        hop_length=hop_length,
        wait=max(1, int(0.03 * sr / hop_length)),
        pre_avg=3,
        post_avg=3,
        pre_max=3,
        post_max=5,
        delta=float(delta),
        backtrack=True,
        units="frames",
    )
    return int(len(frames))


def _auto_tune_params(audio: np.ndarray, sr: int, params: PipelineParams) -> dict[str, Any]:
    """Tune the common detection/grouping parameters from the separated target stem."""
    duration = max(0.001, len(audio) / max(sr, 1))
    # The target density keeps enough repeated hits for clustering without drowning the UI in bleed.
    desired_min = max(4, int(duration * 1.1))
    desired_max = max(desired_min + 2, int(duration * 4.8))
    candidates = [0.06, 0.075, 0.09, 0.105, 0.12, 0.145, 0.17, 0.21, 0.26, 0.32]
    envelope = _build_onset_envelope(audio, sr)
    counts = [(delta, _count_onsets_from_envelope(envelope, sr, delta)) for delta in candidates]
    viable = [(delta, count) for delta, count in counts if desired_min <= count <= desired_max]
    if viable:
        chosen_delta, chosen_count = min(viable, key=lambda item: abs(item[1] - ((desired_min + desired_max) / 2)))
    else:
        chosen_delta, chosen_count = min(counts, key=lambda item: min(abs(item[1] - desired_min), abs(item[1] - desired_max)))
    type_budget = max(6, min(48, int(round((chosen_count or desired_min) ** 0.5 * 3))))
    params.onset_delta = float(chosen_delta)
    params.target_min = max(2, min(int(params.target_min or 0) or 5, type_budget))
    params.target_max = max(params.target_min, type_budget)
    if params.clustering_mode == "batch_quality" and duration <= 180:
        # Keep quality for normal songs; online preview remains user/preset controlled.
        params.clustering_mode = "batch_quality"
    return {
        "chosen_onset_delta": round(float(chosen_delta), 4),
        "estimated_onsets": int(chosen_count),
        "desired_onset_range": [int(desired_min), int(desired_max)],
        "target_min": int(params.target_min),
        "target_max": int(params.target_max),
        "candidate_counts": [{"onset_delta": round(float(delta), 4), "count": int(count)} for delta, count in counts],
    }

def _make_overview(audio: np.ndarray, sr: int, hits: list[Any], max_points: int = 1600) -> dict[str, Any]:
    if len(audio) == 0:
        return {"sample_rate": sr, "duration_sec": 0, "envelope": [], "onsets": []}
    frame = max(1, int(np.ceil(len(audio) / max_points)))
    usable = (len(audio) // frame) * frame
    if usable == 0:
        envelope = [float(np.max(np.abs(audio)))]
    else:
        envelope = np.max(np.abs(audio[:usable].reshape(-1, frame)), axis=1).astype(float).tolist()
    return {
        "sample_rate": sr,
        "duration_sec": round(len(audio) / sr, 6),
        "frame_duration_sec": round(frame / sr, 6),
        "envelope": [round(float(x), 6) for x in envelope],
        "onsets": [
            {
                "index": int(getattr(h, "index", -1)),
                "time_sec": round(float(h.onset_time), 6),
                "duration_sec": round(float(h.duration), 6),
                "label": h.label,
                "energy": round(float(h.rms_energy), 6),
                "cluster_id": int(getattr(h, "cluster_id", -1)),
            }
            for h in hits
        ],
    }


def _copy_temp_file(src: str | os.PathLike[str], dst: Path) -> str:
    dst.parent.mkdir(parents=True, exist_ok=True)
    shutil.copyfile(src, dst)
    return str(dst)


def _safe_file_component(value: str) -> str:
    cleaned = "".join(ch if ch.isalnum() or ch in {"-", "_"} else "_" for ch in value.lower())
    while "__" in cleaned:
        cleaned = cleaned.replace("__", "_")
    return cleaned.strip("_") or "item"


def run_extraction_pipeline(
    audio_path: str | os.PathLike[str],
    output_dir: str | os.PathLike[str],
    params: PipelineParams | dict[str, Any] | None = None,
    progress_cb: ProgressCallback | None = None,
) -> PipelineResult:
    """Run extraction and write all runtime artifacts into output_dir."""
    if not isinstance(params, PipelineParams):
        params = PipelineParams.from_mapping(params)

    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    samples_dir = out / "samples"
    samples_dir.mkdir(parents=True, exist_ok=True)

    stages = [StageTiming(key=key, label=label) for key, label in STAGE_DEFS]
    started_total = time.perf_counter()

    bpm: float | None = None
    stem_audio: np.ndarray
    source_audio: np.ndarray
    context_bed: np.ndarray
    stem_sr: int
    hits: list[Any] = []
    clusters: list[Any] = []
    target_rendered: np.ndarray | None = None
    reproduced: np.ndarray | None = None

    _notify(progress_cb, {"type": "start", "stages": [asdict(s) for s in stages], "progress": _progress_payload(stages)})

    with _timed_stage(stages, "stem", progress_cb) as stage:
        def _stem_progress(event: dict[str, Any]) -> None:
            _update_stage_progress(
                stages,
                "stem",
                progress_cb,
                fraction=float(event.get("fraction", 0.0)),
                detail=str(event.get("detail") or stage.detail or "stem extraction"),
                work_done=event.get("completed_units"),
                work_total=event.get("total_units"),
            )

        raw_stem_audio, stem_sr, stem_detail, separated_context = _load_or_extract_separation(audio_path, params, progress_cb=_stem_progress)
        source_raw = _load_source_mix(audio_path, stem_sr)
        length = max(len(raw_stem_audio), len(source_raw))
        raw_stem_audio = _pad_or_trim(raw_stem_audio, length)
        source_raw = _pad_or_trim(source_raw, length)
        gain = _common_gain(raw_stem_audio if params.stem != "all" else source_raw, source_raw)
        stem_audio = (raw_stem_audio / gain).astype(np.float32)
        source_audio = (source_raw / gain).astype(np.float32)
        context_bed = np.zeros_like(source_audio) if params.stem == "all" else (_pad_or_trim(separated_context, length) / gain if separated_context is not None else (source_audio - stem_audio)).astype(np.float32)
        context_detail = "full mix" if params.stem == "all" else ("separated non-target stems" if separated_context is not None else "residual non-target stems")
        stage.detail = stem_detail + f" · reproduction uses {context_detail}"
        _write_audio(out / "source.wav", _soft_limit(source_audio), stem_sr, subtype="PCM_16")
        _write_audio(out / "stem.wav", _soft_limit(stem_audio), stem_sr, subtype="PCM_16")
        _write_audio(out / "context_bed.wav", _soft_limit(context_bed), stem_sr, subtype="PCM_16")

    audio_duration_sec = len(source_audio) / stem_sr if stem_sr else 0.0

    with _timed_stage(stages, "auto_tune", progress_cb) as stage:
        if params.auto_tune:
            tuning = _auto_tune_params(stem_audio, stem_sr, params)
            stage.detail = (
                f"onset_delta={tuning['chosen_onset_delta']} · "
                f"estimated {tuning['estimated_onsets']} hits · "
                f"groups {tuning['target_min']}-{tuning['target_max']}"
            )
        else:
            stage.detail = "disabled"

    with _timed_stage(stages, "bpm", progress_cb) as stage:
        bpm = detect_bpm(stem_audio, stem_sr)
        stage.detail = f"{bpm} BPM"

    with _timed_stage(stages, "onsets", progress_cb) as stage:
        hits = detect_onsets(
            stem_audio,
            stem_sr,
            mode=params.onset_mode,
            onset_delta=float(params.onset_delta),
            energy_threshold_db=float(params.energy_threshold_db),
            pre_pad=float(params.pre_pad),
            min_dur=float(params.min_dur),
            max_dur=float(params.max_dur),
            min_gap=float(params.min_gap),
        )
        stage.detail = f"{len(hits)} hits"

    if hits:
        with _timed_stage(stages, "classification", progress_cb) as stage:
            hits = classify_hits(hits)
            counts: dict[str, int] = {}
            for hit in hits:
                counts[hit.label] = counts.get(hit.label, 0) + 1
            stage.detail = ", ".join(f"{key}:{value}" for key, value in sorted(counts.items()))

        with _timed_stage(stages, "clustering", progress_cb) as stage:
            if params.clustering_mode == "online_preview":
                clusters = cluster_hits_online(
                    hits,
                    audio=stem_audio,
                    sr=stem_sr,
                    ncc_threshold=float(params.ncc_threshold),
                    attack_ms=float(params.attack_ms),
                    mel_threshold=float(params.mel_threshold),
                    target_min=int(params.target_min),
                    target_max=int(params.target_max),
                )
                stage.detail = f"{len(clusters)} clusters · online preview"
            else:
                clusters = cluster_hits(
                    hits,
                    audio=stem_audio,
                    sr=stem_sr,
                    ncc_threshold=float(params.ncc_threshold),
                    attack_ms=float(params.attack_ms),
                    mel_threshold=float(params.mel_threshold),
                    target_min=int(params.target_min),
                    target_max=int(params.target_max),
                    linkage=params.linkage,
                )
                stage.detail = f"{len(clusters)} clusters · batch quality"
            for cluster in clusters:
                for hit in cluster.hits:
                    hit.cluster_id = cluster.cluster_id

        with _timed_stage(stages, "selection", progress_cb) as stage:
            select_best(clusters)
            stage.detail = "quality-scored representatives"

        with _timed_stage(stages, "synthesis", progress_cb) as stage:
            if params.synthesize:
                synth_count = 0
                for cluster in clusters:
                    if cluster.count >= 2:
                        cluster.synthesized = synthesize_from_cluster(cluster)
                        synth_count += int(cluster.synthesized is not None)
                stage.detail = f"{synth_count} synthesized alternates"
            else:
                stage.detail = "disabled"
    else:
        for key, detail in [
            ("classification", "skipped: no hits"),
            ("clustering", "skipped: no hits"),
            ("selection", "skipped: no hits"),
            ("synthesis", "skipped: no hits"),
        ]:
            stage = next(s for s in stages if s.key == key)
            stage.status = "done"
            stage.detail = detail

    sample_rows: list[dict[str, Any]] = []
    hit_rows: list[dict[str, Any]] = []
    files: dict[str, str] = {"source": "source.wav", "stem": "stem.wav", "context_bed": "context_bed.wav"}

    with _timed_stage(stages, "export", progress_cb) as stage:
        midi_path = out / "reconstruction.mid"
        if clusters:
            export_midi(
                clusters,
                str(midi_path),
                bpm=bpm or 120.0,
                quantize=bool(params.quantize_midi),
                subdivision=int(params.subdivision),
            )
            target_rendered = render_midi_with_samples(clusters, sr=stem_sr)
            target_rendered = _match_rms(target_rendered, stem_audio)
        else:
            target_rendered = np.zeros_like(stem_audio)
            midi_path.write_bytes(b"")

        reproduced = _make_reproduction_mix(target_rendered, context_bed, max(len(source_audio), len(target_rendered)))
        _write_audio(out / "target_reconstruction.wav", _soft_limit(target_rendered), stem_sr, subtype="PCM_16")
        _write_audio(out / "reconstruction.wav", reproduced, stem_sr, subtype="PCM_16")
        files["target_reconstruction"] = "target_reconstruction.wav"
        files["reconstruction"] = "reconstruction.wav"
        files["midi"] = "reconstruction.mid"

        cluster_labels = {int(cluster.cluster_id): cluster.label for cluster in clusters}
        representative_ids = {id(cluster.best_hit) for cluster in clusters}
        review_hits_dir = out / "review" / "hits"
        if hits:
            review_hits_dir.mkdir(parents=True, exist_ok=True)
        for hit in sorted(hits, key=lambda item: item.index):
            safe_label = _safe_file_component(hit.label or "hit")
            file_name = f"hit_{int(hit.index):05d}_{safe_label}.wav"
            rel_file = f"review/hits/{file_name}"
            hit.save(str(out / rel_file))
            cluster_id = int(getattr(hit, "cluster_id", -1))
            hit_rows.append(
                {
                    "index": int(hit.index),
                    "label": hit.label,
                    "cluster_id": cluster_id,
                    "cluster_label": cluster_labels.get(cluster_id, "unclustered"),
                    "is_representative": id(hit) in representative_ids,
                    "onset_sec": round(float(hit.onset_time), 6),
                    "duration_ms": round(float(hit.duration * 1000), 1),
                    "rms_energy": round(float(hit.rms_energy), 6),
                    "spectral_centroid_hz": round(float(hit.spectral_centroid), 1),
                    "file": rel_file,
                }
            )

        for cluster in sorted(clusters, key=lambda item: item.count, reverse=True):
            best = cluster.best_hit
            sample_path = samples_dir / f"{cluster.label}.wav"
            best.save(str(sample_path))
            quality = sample_quality_score(best.audio, best.sr, cluster.label.rsplit("_", 1)[0])
            sample_row = {
                "label": cluster.label,
                "classification": cluster.label.rsplit("_", 1)[0],
                "hits": int(cluster.count),
                "midi_note": int(cluster.midi_note),
                "score": round(float(quality["total"]), 2),
                "cleanness": round(float(quality["cleanness"]), 4),
                "completeness": round(float(quality["completeness"]), 4),
                "duration_ms": round(float(best.duration * 1000), 1),
                "first_onset_sec": round(float(min(hit.onset_time for hit in cluster.hits)), 4),
                "representative_hit_index": int(best.index),
                "cluster_id": int(cluster.cluster_id),
                "file": f"samples/{cluster.label}.wav",
            }
            sample_rows.append(sample_row)
            _notify(progress_cb, {"type": "sample", "sample": sample_row, "samples": list(sample_rows)})
            if cluster.synthesized is not None:
                synth_path = samples_dir / f"{cluster.label}__synth.wav"
                _write_audio(synth_path, cluster.synthesized, stem_sr)

        archive_tmp = build_archive(
            clusters,
            bpm or 120.0,
            stem_sr,
            midi_path=str(midi_path),
            rendered_audio=reproduced,
            target_rendered_audio=target_rendered,
        )
        files["archive"] = _copy_temp_file(archive_tmp, out / "sample-pack.zip")
        files["archive"] = "sample-pack.zip"
        try:
            os.unlink(archive_tmp)
        except OSError:
            pass
        stage.detail = f"{len(sample_rows)} samples + {len(hit_rows)} review hits + MIDI + full-context reproduction + ZIP"

    duration_sec = time.perf_counter() - started_total
    result = PipelineResult(
        params=asdict(params),
        duration_sec=round(duration_sec, 6),
        audio_duration_sec=round(audio_duration_sec, 6),
        realtime_factor=round(duration_sec / max(audio_duration_sec, 1e-9), 6),
        bpm=bpm,
        sample_rate=stem_sr,
        hit_count=len(hits),
        cluster_count=len(clusters),
        stages=[asdict(stage) for stage in stages],
        samples=sample_rows,
        hits=hit_rows,
        overview=_make_overview(stem_audio, stem_sr, hits),
        files=files,
    )

    (out / "manifest.json").write_text(json.dumps(asdict(result), indent=2), encoding="utf-8")
    _notify(progress_cb, {"type": "complete", "result": asdict(result), "stages": result.stages, "progress": {**_progress_payload(stages), "fraction": 1.0}})
    return result