polyvoice 0.20.0

//! polyvoice-bench — DER on a {audio,rttm} dataset directory.
//!
//! Default pipeline matches the shipped CLI (**v2 + VBx** since 0.11). Pass
//! `--pipeline legacy` for the pre-0.11 Silero + AHC path, or `--clusterer ahc`
//! to keep v2 segmentation with fixed-threshold AHC.

use anyhow::{Context, Result};
use clap::Parser;
use polyvoice::cli_common;
use polyvoice::der::{
    DerResult, compute_der, compute_der_decomposition, compute_der_single_speaker_regions,
    compute_der_with_uem, parse_uem,
};
use polyvoice::models::ModelRegistry;
#[cfg(feature = "onnx")]
use polyvoice::pipeline::LegacyPipeline;
use polyvoice::pipeline_v2::{Pipeline as V2Pipeline, PipelineConfig, StageTimings};
use polyvoice::types::{DiarizationResult, Profile, SampleRate, TimeRange};
#[cfg(feature = "onnx")]
use polyvoice::vad::VadConfig;
use polyvoice::wav::read_wav;
use serde::Serialize;
use sha2::{Digest, Sha256};
use std::collections::{HashMap, HashSet, VecDeque};
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;

#[derive(Parser, Debug, Clone)]
#[command(name = "polyvoice-bench", about = "Run DER on a {audio,rttm} dataset")]
struct Args {
    dataset: PathBuf,
    #[arg(long, default_value = "balanced")]
    profile: String,
    #[arg(long)]
    output: Option<PathBuf>,
    #[arg(long, default_value = "0.25")]
    collar: f64,
    /// Score the headline DER over single-speaker reference regions only
    /// (md-eval skip-overlap semantics): overlap frames are excluded from both
    /// the speaker mapping and the error counts, per file and in all
    /// aggregates. Incompatible with --uem.
    #[arg(long, default_value = "false")]
    skip_overlap: bool,
    #[arg(long)]
    max_files: Option<usize>,
    /// How many files to diarize in parallel. Default 1 keeps historical
    /// DER/RTF reports single-threaded. The v2 pipeline is shared by all
    /// workers (jobs>1 adds no model memory); legacy builds a pipeline per
    /// worker. Internal window/embed fan-out is divided by jobs, so jobs ×
    /// workers stays near core count.
    #[arg(long, default_value_t = 1)]
    jobs: usize,
    /// AHC merge threshold on the active scorer's scale: raw cosine (default
    /// 0.45), or AS-norm z-score when `--as-norm` is set (default 4.0).
    /// An explicit value wins over `--domain-profile`'s calibrated threshold.
    #[arg(long)]
    threshold: Option<f32>,
    /// Which pipeline to benchmark: `v2` (powerset segmentation + embeddings +
    /// clusterer + overlap resegmentation — the shipped CLI default) or
    /// `legacy` (Silero VAD + sliding-window embeddings + AHC).
    #[arg(long, default_value = "v2")]
    pipeline: String,
    /// Min cluster size (members): clusters smaller than this are dissolved into
    /// the nearest large speaker. Applies to both pipelines.
    #[arg(long)]
    min_cluster_size: Option<usize>,
    /// v2 clusterer: `vbx` (Variational Bayes HMM + PLDA with automatic speaker
    /// count — matches CLI default; PLDA from env/dir/registry) or `ahc`
    /// (fixed-threshold AHC). Ignored with `--pipeline legacy`.
    #[arg(long, default_value = "vbx")]
    clusterer: String,
    /// Min cluster duration in seconds (length-invariant pruning). When > 0 it
    /// takes precedence over --min-cluster-size on the legacy pipeline.
    #[arg(long)]
    min_cluster_secs: Option<f64>,
    /// Optional .uem file. Restricts DER to the scored regions per file (frames
    /// outside the UEM are dropped from both mapping and counts).
    #[arg(long)]
    uem: Option<PathBuf>,
    /// v2 dense embedding window (seconds): split segments into `w`-sec windows
    /// (hop w/2) for more embeddings per speaker. Omit for one embedding/segment.
    #[arg(long)]
    embed_window: Option<f32>,
    /// ONNX execution provider: auto|cpu|coreml|nnapi|cuda|xnnpack. Omitted =
    /// each pipeline's shipped default (legacy embedder: cpu; v2: auto), so
    /// committed DER baselines stay reproducible. The resolved provider is
    /// recorded in the report for per-backend RTFx comparison.
    #[arg(long)]
    execution_provider: Option<String>,
    /// v2 binarization: enter-speech (onset) threshold. Setting ANY
    /// --binarize-* flag enables calibrated hysteresis binarization of the
    /// segmentation posteriors (defaults for unset knobs: onset/offset 0.5,
    /// min durations 0).
    #[arg(long)]
    binarize_onset: Option<f32>,
    /// v2 binarization: leave-speech (offset) threshold (< onset = hysteresis).
    #[arg(long)]
    binarize_offset: Option<f32>,
    /// v2 binarization: drop active runs shorter than this many seconds.
    #[arg(long)]
    binarize_min_on: Option<f32>,
    /// v2 binarization: bridge gaps shorter than this many seconds.
    #[arg(long)]
    binarize_min_off: Option<f32>,
    /// AS-norm score normalization for the AHC clusterer (v2 only; requires
    /// --clusterer ahc): pairwise cosine scores are z-normalized against an
    /// imposter cohort before merging, so one threshold generalizes across
    /// domains.
    #[arg(long)]
    as_norm: bool,
    /// Imposter cohort for --as-norm: (N, 256) '<f4' .npy of speaker
    /// embeddings. Omitted = model-registry cohort (id asnorm_cohort_voxdev),
    /// with the POLYVOICE_ASNORM_COHORT env override.
    #[arg(long)]
    cohort: Option<PathBuf>,
    /// Per-domain scoring profile: voxconverse | ami | callhome. Replaces the
    /// default AHC threshold with the profile's calibrated value (an explicit
    /// --threshold wins) and sets the AS-norm cohort size. Requires
    /// --clusterer ahc. v2 only.
    #[arg(long)]
    domain_profile: Option<String>,
}

#[derive(Serialize)]
struct ModelHash {
    model_id: String,
    sha256: String,
}

#[derive(Serialize)]
struct PerSpeakerRecall {
    speaker: u32,
    recall: f64,
}

#[derive(Serialize)]
struct PerFileResult {
    filename: String,
    der_collar: f64,
    der_no_collar: f64,
    miss_rate: f64,
    false_alarm_rate: f64,
    confusion_rate: f64,
    /// Overlap-aware decomposition: DER over single-speaker reference regions
    /// only, DER over overlap regions only (>= 2 ref speakers), and per-speaker
    /// recall. All at the requested collar. Makes overlap-heavy DER interpretable.
    der_single_speaker: f64,
    der_overlap: f64,
    per_speaker_recall: Vec<PerSpeakerRecall>,
    rt_factor: f64,
    ref_speakers: usize,
    hyp_speakers: usize,
    num_turns: usize,
    audio_duration_secs: f64,
    runtime_secs: f64,
    /// Per-stage wall-clock seconds (v2 pipeline only; absent on legacy).
    #[serde(skip_serializing_if = "Option::is_none")]
    stage_timings: Option<StageTimings>,
}

#[derive(Serialize)]
struct SpeakerCountDiagnostics {
    exact: usize,
    plus_minus_1: usize,
    off_by_2_or_more: usize,
}

#[derive(Serialize)]
struct BenchReport {
    schema: &'static str,
    crate_version: &'static str,
    git_sha: String,
    host_arch: String,
    host_os: String,
    command_line: String,
    dataset_name: String,
    profile: String,
    files_processed: usize,
    files_skipped: usize,
    /// Mean of per-file DER (macro) at the requested collar and at collar=0.
    der_collar_macro: f64,
    der_no_collar_macro: f64,
    /// Duration-weighted DER (micro): sum of error frames / sum of reference
    /// frames — comparable to pyannote/speakrs headline numbers.
    der_collar_micro: f64,
    der_no_collar_micro: f64,
    collar_secs: f64,
    /// True when --skip-overlap was active: the headline DER (per file and in
    /// all aggregates) is computed over single-speaker reference regions only.
    skip_overlap: bool,
    averaging_policy: &'static str,
    /// Debug-formatted resolved execution provider (e.g. "CoreMl", "Cpu") —
    /// labels every report for per-backend RTFx comparison.
    resolved_execution_provider: String,
    host_cpus: usize,
    /// Sum of per-stage wall-clock seconds across files (v2 only).
    #[serde(skip_serializing_if = "Option::is_none")]
    stage_totals: Option<StageTimings>,
    miss: f64,
    false_alarm: f64,
    confusion: f64,
    rt_factor_avg: f64,
    /// Total audio seconds / wall-clock seconds of the file loop. Unlike
    /// `rt_factor_avg` (sum of per-file runtimes) this reflects file-level
    /// parallelism when `--jobs > 1`.
    rt_factor_wall: f64,
    speaker_count: SpeakerCountDiagnostics,
    model_hashes: Vec<ModelHash>,
    per_file: Vec<PerFileResult>,
}

fn git_sha() -> String {
    std::process::Command::new("git")
        .args(["rev-parse", "HEAD"])
        .output()
        .ok()
        .and_then(|o| {
            if o.status.success() {
                String::from_utf8(o.stdout).ok()
            } else {
                None
            }
        })
        .map(|s| s.trim().to_owned())
        .unwrap_or_else(|| "unknown".to_owned())
}

fn model_hashes(registry: &ModelRegistry, profile: Profile, segmenter_id: &str) -> Vec<ModelHash> {
    let mut out = Vec::new();
    let manifest = registry.manifest();
    let prof = match manifest.profile(profile.manifest_id()) {
        Some(p) => p,
        None => return out,
    };
    // Report exactly the models the chosen pipeline actually loads: the legacy
    // path segments with Silero VAD, the v2 path with the profile's powerset
    // segmenter — `segmenter_id` carries the right one. Both embed with the
    // profile embedder. This keeps the integrity record honest about what
    // produced the DER number.
    for model_id in [segmenter_id, prof.embedder.as_str()] {
        if let Some(entry) = manifest.model(model_id) {
            out.push(ModelHash {
                model_id: model_id.to_string(),
                sha256: entry.sha256.clone(),
            });
        }
    }
    out
}

/// Hard-fail unless the on-disk embedder + VAD match the manifest sha256, so a DER
/// number can never be silently attributed to a swapped/corrupted/non-FP32 model.
#[cfg(feature = "onnx")]
fn verify_model_integrity(
    registry: &ModelRegistry,
    profile: Profile,
    embedder_path: &Path,
    vad_path: &Path,
) -> Result<()> {
    let manifest = registry.manifest();
    let prof = manifest
        .profile(profile.manifest_id())
        .ok_or_else(|| anyhow::anyhow!("profile {} not in manifest", profile.manifest_id()))?;
    check_model_sha256(registry, &prof.embedder, embedder_path)?;
    check_model_sha256(registry, "silero_vad", vad_path)?;
    Ok(())
}

fn check_model_sha256(registry: &ModelRegistry, model_id: &str, path: &Path) -> Result<()> {
    let manifest = registry.manifest();
    let entry = manifest
        .model(model_id)
        .ok_or_else(|| anyhow::anyhow!("model {model_id} not in manifest"))?;
    let bytes = std::fs::read(path).with_context(|| format!("read model {}", path.display()))?;
    let got = hex_lower(&Sha256::digest(&bytes));
    if !got.eq_ignore_ascii_case(&entry.sha256) {
        anyhow::bail!(
            "model integrity FAIL for {model_id}: on-disk sha256 {got} != manifest {}",
            entry.sha256
        );
    }
    Ok(())
}

fn hex_lower(bytes: &[u8]) -> String {
    use std::fmt::Write as _;
    let mut s = String::with_capacity(bytes.len() * 2);
    for b in bytes {
        let _ = write!(s, "{b:02x}");
    }
    s
}

/// Legacy pipeline + its ONNX sessions (Silero VAD + sliding-window embedder).
#[cfg(feature = "onnx")]
struct LegacyRunner {
    pipeline: LegacyPipeline,
    stack: cli_common::LegacyStack,
}

/// The pipeline under benchmark. Both arms produce a `DiarizationResult` so all
/// downstream DER / speaker-count reporting is shared. Both payloads are boxed
/// so the variants are the same (pointer) size.
enum Runner {
    #[cfg(feature = "onnx")]
    Legacy(Box<LegacyRunner>),
    V2(Box<V2Pipeline>),
}

impl Runner {
    fn run(
        &mut self,
        samples: &[f32],
        sr: SampleRate,
    ) -> Result<(DiarizationResult, Option<StageTimings>)> {
        match self {
            #[cfg(feature = "onnx")]
            Runner::Legacy(l) => Ok((
                l.pipeline
                    .run(samples, &l.stack.extractor, &mut l.stack.vad)?,
                None,
            )),
            Runner::V2(p) => {
                let (result, timings) = p.run_with_timings(samples, sr)?;
                Ok((result, Some(timings)))
            }
        }
    }
}

/// The runner plus everything the report needs that is file-independent.
struct BenchRunner {
    runner: Runner,
    segmenter_id: String,
    resolved_ep: polyvoice::pipeline_v2::ExecutionProvider,
    profile: Profile,
    registry: ModelRegistry,
}

/// Build the requested pipeline. Each arm verifies the integrity of exactly
/// the models it loads and yields the segmenter id for the report.
fn build_runner(args: &Args) -> Result<BenchRunner> {
    if args.pipeline == "legacy" {
        cli_common::require_onnx("--pipeline legacy")?;
    }
    let profile: Profile = args.profile.parse()?;
    let registry = ModelRegistry::default().context("registry")?;
    let models = registry
        .ensure_for_profile(profile)
        .context("ensure models")?;

    // Resolve the execution provider: an explicit flag applies to the selected
    // pipeline; omitted keeps each pipeline's shipped default (legacy embedder
    // cpu, v2 auto) so committed DER baselines stay reproducible.
    let explicit_ep = args
        .execution_provider
        .as_deref()
        .map(cli_common::parse_execution_provider)
        .transpose()?;
    let resolved_ep = match args.pipeline.as_str() {
        "v2" => explicit_ep.unwrap_or_else(polyvoice::pipeline_v2::ExecutionProvider::auto),
        _ => explicit_ep.unwrap_or(polyvoice::pipeline_v2::ExecutionProvider::Cpu),
    };

    let (runner, segmenter_id): (Runner, String) = match args.pipeline.as_str() {
        "v2" => {
            let (clusterer, as_norm, domain) = cli_common::resolve_clusterer_flags(
                &args.clusterer,
                args.threshold,
                args.as_norm,
                args.cohort.clone(),
                args.domain_profile.as_deref(),
            )?;
            let binarization = if args.binarize_onset.is_some()
                || args.binarize_offset.is_some()
                || args.binarize_min_on.is_some()
                || args.binarize_min_off.is_some()
            {
                let d = polyvoice::segmentation::BinarizationConfig::default();
                Some(polyvoice::segmentation::BinarizationConfig {
                    onset: args.binarize_onset.unwrap_or(d.onset),
                    offset: args.binarize_offset.unwrap_or(d.offset),
                    min_duration_on: args.binarize_min_on.unwrap_or(d.min_duration_on),
                    min_duration_off: args.binarize_min_off.unwrap_or(d.min_duration_off),
                })
            } else {
                None
            };
            let mut cfg = PipelineConfig {
                profile,
                clusterer,
                embed_window_secs: args.embed_window,
                execution_provider: resolved_ep,
                binarization,
                as_norm,
                domain,
                ..PipelineConfig::default()
            };
            if let Some(mcs) = args.min_cluster_size {
                cfg.min_cluster_size = mcs;
            }
            // v2 segments with the profile's powerset model — verify it + embedder.
            let seg_id = registry
                .manifest()
                .profile(profile.manifest_id())
                .map(|p| p.segmenter.clone())
                .unwrap_or_else(|| "powerset_fp32".to_owned());
            let emb_id = registry
                .manifest()
                .profile(profile.manifest_id())
                .map(|p| p.embedder.clone())
                .unwrap_or_default();
            check_model_sha256(&registry, &seg_id, &models.segmenter_path)?;
            check_model_sha256(&registry, &emb_id, &models.embedder_path)?;
            let pipeline = cli_common::build_v2_pipeline(cfg, registry.clone())?;
            (Runner::V2(Box::new(pipeline)), seg_id)
        }
        other => {
            if other != "legacy" {
                anyhow::bail!("unknown --pipeline '{other}' (expected 'legacy' or 'v2')");
            }
            if args.as_norm || args.cohort.is_some() || args.domain_profile.is_some() {
                anyhow::bail!("--as-norm/--cohort/--domain-profile apply to --pipeline v2 only");
            }
            #[cfg(not(feature = "onnx"))]
            {
                let _ = (models, resolved_ep, registry, profile);
                unreachable!("require_onnx rejected --pipeline legacy");
            }
            #[cfg(feature = "onnx")]
            {
                let vad_path = registry.ensure("silero_vad").context("silero_vad model")?;
                let stack = cli_common::load_legacy_stack(
                    &models.embedder_path,
                    profile.embedding_dim(),
                    resolved_ep,
                    &vad_path,
                    512,
                )?;
                verify_model_integrity(&registry, profile, &models.embedder_path, &vad_path)?;
                let mut config = cli_common::legacy_diarization_config(
                    args.threshold.unwrap_or(polyvoice::DEFAULT_AHC_THRESHOLD),
                );
                config.cluster.min_cluster_size = args.min_cluster_size.unwrap_or(1);
                config.cluster.min_cluster_secs = args.min_cluster_secs.unwrap_or(0.0);
                let pipeline = LegacyPipeline::new(config, VadConfig::default());
                (
                    Runner::Legacy(Box::new(LegacyRunner { pipeline, stack })),
                    "silero_vad".to_owned(),
                )
            }
        }
    };
    Ok(BenchRunner {
        runner,
        segmenter_id,
        resolved_ep,
        profile,
        registry,
    })
}

/// One scored file: the report row plus the inputs the aggregates need.
struct FileOutcome {
    row: PerFileResult,
    der_pair: (DerResult, DerResult),
    ref_count: usize,
    hyp_count: usize,
    audio_secs: f64,
    runtime_secs: f64,
}

/// Process `wavs` serially (`jobs == 1`) or with file-level parallelism. The
/// v2 pipeline is shared by all workers (every stage is `&self` and the
/// traits are `Send + Sync`), so `jobs > 1` costs no extra model memory; the
/// legacy path needs `&mut` state and keeps one runner per worker. Internal
/// fan-out (windows, embed pool) is divided by `jobs` via the kernels
/// file-parallelism hint.
fn run_all_files(
    args: &Args,
    first: &mut BenchRunner,
    wavs: &[PathBuf],
    rttm_dir: &Path,
    uem_map: Option<&HashMap<String, Vec<TimeRange>>>,
) -> Result<Accum> {
    let jobs = args.jobs.max(1).min(wavs.len().max(1));
    #[cfg(feature = "segmenter-native")]
    polyvoice_kernels::set_file_parallelism(jobs);
    if jobs == 1 {
        let wall = Instant::now();
        let mut acc = Accum::default();
        let mut run = |s: &[f32], sr: SampleRate| first.runner.run(s, sr);
        for wav in wavs {
            match run_file(&mut run, wav, rttm_dir, uem_map, args)? {
                Some(outcome) => acc.record(outcome),
                None => acc.files_skipped += 1,
            }
        }
        acc.wall_secs = wall.elapsed().as_secs_f64();
        return Ok(acc);
    }

    match &mut first.runner {
        Runner::V2(p) => run_v2_shared(p, args, wavs, rttm_dir, uem_map, jobs),
        #[cfg(feature = "onnx")]
        Runner::Legacy(_) => run_legacy_parallel(first, args, wavs, rttm_dir, uem_map, jobs),
    }
}

/// v2 file fan-out on one shared pipeline: workers pull paths off a queue,
/// results are collected under a mutex and sorted by filename, so the report
/// is identical to the serial run.
fn run_v2_shared(
    pipeline: &V2Pipeline,
    args: &Args,
    wavs: &[PathBuf],
    rttm_dir: &Path,
    uem_map: Option<&HashMap<String, Vec<TimeRange>>>,
    jobs: usize,
) -> Result<Accum> {
    eprintln!("file-parallel jobs={jobs} (shared pipeline)");
    let queue = Mutex::new(wavs.iter().cloned().collect::<VecDeque<_>>());
    let collected = Mutex::new(Vec::new());
    let skipped = AtomicUsize::new(0);
    let err = Mutex::new(None::<String>);

    let drain = || {
        let mut run = |s: &[f32], sr: SampleRate| {
            pipeline
                .run_with_timings(s, sr)
                .map(|(r, t)| (r, Some(t)))
                .map_err(anyhow::Error::from)
        };
        loop {
            if err
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .is_some()
            {
                return;
            }
            let next = queue
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .pop_front();
            let Some(wav) = next else {
                return;
            };
            match run_file(&mut run, &wav, rttm_dir, uem_map, args) {
                Ok(Some(outcome)) => collected
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner)
                    .push(outcome),
                Ok(None) => {
                    skipped.fetch_add(1, Ordering::Relaxed);
                }
                Err(e) => {
                    *err.lock()
                        .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(e.to_string());
                    return;
                }
            }
        }
    };

    let wall = Instant::now();
    // Spawn the same `Fn` closure `jobs` times: a shared reference is
    // `FnOnce + Send`, so no per-spawn wrapper closure is needed.
    let drain_ref = &drain;
    std::thread::scope(|s| {
        for _ in 0..jobs {
            s.spawn(drain_ref);
        }
    });
    let wall_secs = wall.elapsed().as_secs_f64();

    if let Some(msg) = err
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
        .take()
    {
        anyhow::bail!("{msg}");
    }

    let mut rows = collected
        .into_inner()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    rows.sort_by(|a, b| a.row.filename.cmp(&b.row.filename));
    let mut acc = Accum {
        files_skipped: skipped.load(Ordering::Relaxed),
        wall_secs,
        ..Accum::default()
    };
    for outcome in rows {
        acc.record(outcome);
    }
    Ok(acc)
}

/// Legacy file fan-out: `LegacyPipeline::run` needs `&mut` (VAD state), so
/// each worker builds its own runner — models are loaded `jobs` times.
#[cfg(feature = "onnx")]
fn run_legacy_parallel(
    first: &mut BenchRunner,
    args: &Args,
    wavs: &[PathBuf],
    rttm_dir: &Path,
    uem_map: Option<&HashMap<String, Vec<TimeRange>>>,
    jobs: usize,
) -> Result<Accum> {
    eprintln!("file-parallel jobs={jobs} (legacy: one pipeline per worker)");
    let mut extra: Vec<BenchRunner> = Vec::with_capacity(jobs.saturating_sub(1));
    for _ in 1..jobs {
        extra.push(build_runner(args)?);
    }

    let queue = Mutex::new(wavs.iter().cloned().collect::<VecDeque<_>>());
    let collected = Mutex::new(Vec::new());
    let skipped = AtomicUsize::new(0);
    let err = Mutex::new(None::<String>);

    let drain = |runner: &mut Runner| {
        let mut run = |s: &[f32], sr: SampleRate| runner.run(s, sr);
        loop {
            if err
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .is_some()
            {
                return;
            }
            let next = queue
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .pop_front();
            let Some(wav) = next else {
                return;
            };
            match run_file(&mut run, &wav, rttm_dir, uem_map, args) {
                Ok(Some(outcome)) => collected
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner)
                    .push(outcome),
                Ok(None) => {
                    skipped.fetch_add(1, Ordering::Relaxed);
                }
                Err(e) => {
                    *err.lock()
                        .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(e.to_string());
                    return;
                }
            }
        }
    };

    let wall = Instant::now();
    std::thread::scope(|s| {
        s.spawn(|| drain(&mut first.runner));
        for br in &mut extra {
            s.spawn(|| drain(&mut br.runner));
        }
    });
    let wall_secs = wall.elapsed().as_secs_f64();

    if let Some(msg) = err
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
        .take()
    {
        anyhow::bail!("{msg}");
    }

    let mut rows = collected
        .into_inner()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    rows.sort_by(|a, b| a.row.filename.cmp(&b.row.filename));
    let mut acc = Accum {
        files_skipped: skipped.load(Ordering::Relaxed),
        wall_secs,
        ..Accum::default()
    };
    for outcome in rows {
        acc.record(outcome);
    }
    Ok(acc)
}

/// Run one wav through the pipeline and score it. `run` maps PCM to the
/// diarization result (plus v2 stage timings); callers decide whether it
/// borrows a shared pipeline or a per-worker runner. `Ok(None)` means the
/// file was skipped (no reference RTTM).
fn run_file<F>(
    run: &mut F,
    wav: &Path,
    rttm_dir: &Path,
    uem_map: Option<&HashMap<String, Vec<TimeRange>>>,
    args: &Args,
) -> Result<Option<FileOutcome>>
where
    F: FnMut(&[f32], SampleRate) -> Result<(DiarizationResult, Option<StageTimings>)>,
{
    let stem = wav.file_stem().and_then(|s| s.to_str()).unwrap_or("");
    let rttm = rttm_dir.join(format!("{stem}.rttm"));
    if !rttm.is_file() {
        eprintln!("[SKIP] {stem}: no rttm");
        return Ok(None);
    }
    let (samples, sr_hz) = read_wav(wav)?;
    let sr =
        SampleRate::new(sr_hz).ok_or_else(|| anyhow::anyhow!("invalid sample rate: {sr_hz}"))?;
    let audio_secs = samples.len() as f64 / sr_hz as f64;

    let t0 = Instant::now();
    let (result, stage_timings) = run(&samples, sr)?;
    let runtime_secs = t0.elapsed().as_secs_f64();

    let ref_turns = cli_common::load_ref_turns(rttm_dir, stem)?;

    // Headline collar + no-collar DER, restricted to the UEM scope when present
    // (AMI-style id fallback like the RTTM lookup). With --skip-overlap the
    // headline is the single-speaker-regions DER (md-eval skip-overlap
    // semantics); that mode rejects --uem at startup. The overlap
    // decomposition below is a diagnostic and stays over the full file.
    let scored: Option<&[TimeRange]> = uem_map.and_then(|m| {
        m.get(stem)
            .or_else(|| stem.split('.').next().and_then(|s| m.get(s)))
            .map(|v| v.as_slice())
    });
    let (der, der_no_collar) = if args.skip_overlap {
        (
            compute_der_single_speaker_regions(&ref_turns, &result.turns, args.collar),
            compute_der_single_speaker_regions(&ref_turns, &result.turns, 0.0),
        )
    } else {
        match scored {
            Some(s) => (
                compute_der_with_uem(&ref_turns, &result.turns, args.collar, s),
                compute_der_with_uem(&ref_turns, &result.turns, 0.0, s),
            ),
            None => (
                compute_der(&ref_turns, &result.turns, args.collar),
                compute_der(&ref_turns, &result.turns, 0.0),
            ),
        }
    };
    let decomp = compute_der_decomposition(&ref_turns, &result.turns, args.collar);

    let ref_speakers: HashSet<_> = ref_turns.iter().map(|t| t.speaker.0).collect();
    let hyp_speakers: HashSet<_> = result.turns.iter().map(|t| t.speaker.0).collect();
    let ref_count = ref_speakers.len();
    let hyp_count = hyp_speakers.len();

    let rt_factor = audio_secs / runtime_secs.max(1e-6);

    println!(
        "{stem}\t DER={:.3}%\t miss={:.3}%\t fa={:.3}%\t conf={:.3}%\t rt={:.1}x\t spk={}\t turns={}",
        der.der * 100.0,
        der.miss_rate * 100.0,
        der.false_alarm_rate * 100.0,
        der.confusion_rate * 100.0,
        rt_factor,
        result.num_speakers,
        result.turns.len(),
    );

    let row = PerFileResult {
        filename: stem.to_owned(),
        der_collar: der.der * 100.0,
        der_no_collar: der_no_collar.der * 100.0,
        miss_rate: der.miss_rate * 100.0,
        false_alarm_rate: der.false_alarm_rate * 100.0,
        confusion_rate: der.confusion_rate * 100.0,
        der_single_speaker: decomp.single_speaker.der * 100.0,
        der_overlap: decomp.overlap.der * 100.0,
        per_speaker_recall: decomp
            .per_speaker_recall
            .iter()
            .map(|s| PerSpeakerRecall {
                speaker: s.speaker,
                recall: s.recall,
            })
            .collect(),
        rt_factor,
        ref_speakers: ref_count,
        hyp_speakers: hyp_count,
        num_turns: result.turns.len(),
        audio_duration_secs: audio_secs,
        runtime_secs,
        stage_timings,
    };
    Ok(Some(FileOutcome {
        row,
        der_pair: (der, der_no_collar),
        ref_count,
        hyp_count,
        audio_secs,
        runtime_secs,
    }))
}

/// Per-run accumulators, folded one [`FileOutcome`] at a time.
#[derive(Default)]
struct Accum {
    totals: Aggregate,
    /// Per-file (collar, no-collar) DER pairs — the four report aggregates are
    /// computed from these by the unit-tested aggregate_der helper.
    der_pairs: Vec<(DerResult, DerResult)>,
    speaker_exact: usize,
    speaker_pm1: usize,
    speaker_off: usize,
    files_skipped: usize,
    total_audio_secs: f64,
    total_runtime_secs: f64,
    /// Wall-clock seconds of the file loop (parallel jobs overlap, so this is
    /// less than `total_runtime_secs` when jobs > 1).
    wall_secs: f64,
    stage_totals: Option<StageTimings>,
    per_file: Vec<PerFileResult>,
}

impl Accum {
    fn record(&mut self, outcome: FileOutcome) {
        let FileOutcome {
            row,
            der_pair,
            ref_count,
            hyp_count,
            audio_secs,
            runtime_secs,
        } = outcome;
        let (der, der_no_collar) = der_pair;
        self.totals.miss += der.miss_rate;
        self.totals.false_alarm += der.false_alarm_rate;
        self.totals.confusion += der.confusion_rate;
        self.totals.count += 1;
        match ref_count.abs_diff(hyp_count) {
            0 => self.speaker_exact += 1,
            1 => self.speaker_pm1 += 1,
            _ => self.speaker_off += 1,
        }
        self.der_pairs.push((der, der_no_collar));
        self.total_audio_secs += audio_secs;
        self.total_runtime_secs += runtime_secs;
        if let Some(t) = &row.stage_timings {
            let acc = self.stage_totals.get_or_insert_with(StageTimings::default);
            acc.segmentation_secs += t.segmentation_secs;
            acc.embedding_secs += t.embedding_secs;
            acc.clustering_secs += t.clustering_secs;
            acc.resegmentation_secs += t.resegmentation_secs;
        }
        self.per_file.push(row);
    }
}

/// Print the aggregate summary and assemble the JSON report. Takes the
/// file-independent runner state as separate fields (not the whole
/// [`BenchRunner`]) so the report assembly is unit-testable without building
/// ONNX pipelines.
#[allow(clippy::too_many_arguments)]
fn build_report(
    args: &Args,
    registry: &ModelRegistry,
    profile: Profile,
    segmenter_id: &str,
    resolved_ep: polyvoice::pipeline_v2::ExecutionProvider,
    dataset_name: String,
    acc: Accum,
) -> BenchReport {
    let n = acc.totals.count.max(1) as f64;
    let agg = aggregate_der(&acc.der_pairs);

    println!(
        "\n=== Aggregate DER over {} files (collar={:.2}s) ===",
        acc.totals.count, args.collar
    );
    if args.skip_overlap {
        println!("  skip-overlap  : ON (single-speaker reference regions only)");
    }
    println!(
        "  der_collar    : macro={:.2}%  micro={:.2}%",
        agg.collar_macro, agg.collar_micro
    );
    println!(
        "  der_no_collar : macro={:.2}%  micro={:.2}%",
        agg.no_collar_macro, agg.no_collar_micro
    );

    BenchReport {
        schema: "polyvoice-bench-v0.10",
        crate_version: env!("CARGO_PKG_VERSION"),
        git_sha: git_sha(),
        host_arch: std::env::consts::ARCH.to_owned(),
        host_os: std::env::consts::OS.to_owned(),
        command_line: std::env::args().collect::<Vec<_>>().join(" "),
        dataset_name,
        profile: args.profile.clone(),
        files_processed: acc.totals.count,
        files_skipped: acc.files_skipped,
        der_collar_macro: agg.collar_macro,
        der_no_collar_macro: agg.no_collar_macro,
        der_collar_micro: agg.collar_micro,
        der_no_collar_micro: agg.no_collar_micro,
        collar_secs: args.collar,
        skip_overlap: args.skip_overlap,
        averaging_policy: "macro = mean of per-file DER; micro = frame-weighted (sum error frames / sum ref frames)",
        resolved_execution_provider: format!("{:?}", resolved_ep),
        host_cpus: std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(1),
        stage_totals: acc.stage_totals,
        miss: (acc.totals.miss / n) * 100.0,
        false_alarm: (acc.totals.false_alarm / n) * 100.0,
        confusion: (acc.totals.confusion / n) * 100.0,
        rt_factor_avg: acc.total_audio_secs / acc.total_runtime_secs.max(1e-6),
        rt_factor_wall: acc.total_audio_secs / acc.wall_secs.max(1e-6),
        speaker_count: SpeakerCountDiagnostics {
            exact: acc.speaker_exact,
            plus_minus_1: acc.speaker_pm1,
            off_by_2_or_more: acc.speaker_off,
        },
        model_hashes: model_hashes(registry, profile, segmenter_id),
        per_file: acc.per_file,
    }
}

fn main() -> Result<()> {
    cli_common::limit_malloc_arenas();
    let args = Args::parse();
    // The DER library has no single-speaker-regions + UEM scorer, so this
    // combination cannot be honoured — fail loudly instead of scoring the
    // wrong thing.
    if args.skip_overlap && args.uem.is_some() {
        anyhow::bail!("--skip-overlap cannot be combined with --uem");
    }
    if args.skip_overlap {
        println!("skip-overlap: headline DER over single-speaker reference regions only");
    }
    let mut b = build_runner(&args)?;

    // Optional UEM scoped regions, keyed by file id.
    let uem_map: Option<HashMap<String, Vec<TimeRange>>> = match &args.uem {
        Some(path) => {
            let text = std::fs::read_to_string(path)
                .with_context(|| format!("read uem {}", path.display()))?;
            Some(parse_uem(&text))
        }
        None => None,
    };

    let wavs = cli_common::list_wavs(&args.dataset, args.max_files)?;
    let rttm_dir = args.dataset.join("rttm");
    let dataset_name = args
        .dataset
        .file_name()
        .and_then(|s| s.to_str())
        .unwrap_or("unknown")
        .to_owned();

    let acc = run_all_files(&args, &mut b, &wavs, &rttm_dir, uem_map.as_ref())?;

    let report = build_report(
        &args,
        &b.registry,
        b.profile,
        &b.segmenter_id,
        b.resolved_ep,
        dataset_name,
        acc,
    );
    let json = serde_json::to_string_pretty(&report)?;
    match args.output {
        Some(p) => std::fs::write(&p, json)?,
        None => println!("{json}"),
    }
    Ok(())
}

#[derive(Default)]
struct Aggregate {
    miss: f64,
    false_alarm: f64,
    confusion: f64,
    count: usize,
}

/// The four report aggregates, as percentages.
struct DerAggregates {
    collar_macro: f64,
    no_collar_macro: f64,
    collar_micro: f64,
    no_collar_micro: f64,
}

/// Compute collar/no-collar x macro/micro DER from per-file result pairs.
/// Macro = mean of per-file ratios; micro = duration-weighted (summed error
/// frames / summed reference frames), with collar and no-collar frame sums
/// kept strictly separate. This is THE aggregation the report publishes —
/// unit-tested so a refactor cannot silently revert micro to a ratio-average
/// or swap the collar passes.
fn aggregate_der(pairs: &[(DerResult, DerResult)]) -> DerAggregates {
    let n = pairs.len().max(1) as f64;
    let (mut cm, mut cf, mut cc, mut cr) = (0u64, 0u64, 0u64, 0u64);
    let (mut nm, mut nf, mut nc, mut nr) = (0u64, 0u64, 0u64, 0u64);
    for (c, n_) in pairs {
        cm += c.missed_frames;
        cf += c.false_alarm_frames;
        cc += c.confusion_frames;
        cr += c.total_ref_frames;
        nm += n_.missed_frames;
        nf += n_.false_alarm_frames;
        nc += n_.confusion_frames;
        nr += n_.total_ref_frames;
    }
    DerAggregates {
        collar_macro: pairs.iter().map(|(c, _)| c.der).sum::<f64>() / n * 100.0,
        no_collar_macro: pairs.iter().map(|(_, x)| x.der).sum::<f64>() / n * 100.0,
        collar_micro: micro_der(cm, cf, cc, cr),
        no_collar_micro: micro_der(nm, nf, nc, nr),
    }
}

/// Duration-weighted micro-average DER as a percentage: total error frames over
/// total reference frames (not a mean of per-file ratios). Returns 0.0 when no
/// reference frames were seen.
fn micro_der(missed: u64, false_alarm: u64, confusion: u64, ref_frames: u64) -> f64 {
    if ref_frames == 0 {
        0.0
    } else {
        (missed + false_alarm + confusion) as f64 / ref_frames as f64 * 100.0
    }
}

#[allow(clippy::unwrap_used)]
#[cfg(test)]
#[path = "polyvoice_bench_prop_tests.rs"]
mod prop_tests;

#[allow(clippy::unwrap_used)]
#[cfg(test)]
#[path = "polyvoice_bench_tests.rs"]
mod tests;