polyvoice 0.10.0

#![allow(deprecated)] // legacy embedding API; see polyvoice::embedder
//! polyvoice-bench — DER on a {audio,rttm} dataset directory using the legacy v0.5 Pipeline.

use anyhow::{Context, Result};
use clap::Parser;
use polyvoice::der::{
    DerResult, compute_der, compute_der_decomposition, compute_der_with_uem, parse_uem,
};
use polyvoice::models::ModelRegistry;
use polyvoice::pipeline::Pipeline;
use polyvoice::pipeline_v2::{ClustererKind, Pipeline as V2Pipeline, PipelineConfig, StageTimings};
use polyvoice::rttm::{group_by_file, parse_rttm_file, to_speaker_turns};
use polyvoice::types::{
    ClusterConfig, DiarizationConfig, DiarizationResult, Profile, SampleRate, TimeRange,
};
use polyvoice::vad::VadConfig;
use polyvoice::wav::read_wav;
use polyvoice::{FbankOnnxExtractor, SileroVad};
use serde::Serialize;
use sha2::{Digest, Sha256};
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::time::Instant;

#[derive(Parser, Debug)]
#[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,
    #[arg(long, default_value = "false")]
    skip_overlap: bool,
    #[arg(long)]
    max_files: Option<usize>,
    #[arg(long, default_value = "0.45")]
    threshold: f32,
    /// Which pipeline to benchmark: `legacy` (Silero VAD + sliding-window
    /// embeddings + AHC, the shipped default) or `v2` (powerset segmentation +
    /// overlap-add + masked embeddings + centroid-AHC + min_cluster_size).
    #[arg(long, default_value = "legacy")]
    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: `ahc` (fixed-threshold AHC, default) or `vbx` (Variational
    /// Bayes HMM + PLDA with automatic speaker count; requires a `vbx` build and
    /// `POLYVOICE_VBX_PLDA_DIR`).
    #[arg(long, default_value = "ahc")]
    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>,
}

#[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,
    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,
    speaker_count: SpeakerCountDiagnostics,
    model_hashes: Vec<ModelHash>,
    per_file: Vec<PerFileResult>,
}

fn parse_profile(name: &str) -> Result<Profile> {
    match name {
        "mobile" => Ok(Profile::Mobile),
        "balanced" => Ok(Profile::Balanced),
        other => anyhow::bail!("invalid profile: {other}"),
    }
}

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.
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).
struct LegacyRunner {
    pipeline: Pipeline,
    extractor: FbankOnnxExtractor,
    vad: SileroVad,
}

/// 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 {
    Legacy(Box<LegacyRunner>),
    V2(Box<V2Pipeline>),
}

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

fn parse_execution_provider(s: &str) -> Result<polyvoice::onnx::ExecutionProvider> {
    use polyvoice::onnx::ExecutionProvider as Ep;
    Ok(match s {
        "auto" => Ep::auto(),
        "cpu" => Ep::Cpu,
        "coreml" => Ep::CoreMl,
        "nnapi" => Ep::Nnapi,
        "cuda" => Ep::Cuda,
        "xnnpack" => Ep::XnnPack,
        other => anyhow::bail!(
            "unknown --execution-provider '{other}' (expected auto|cpu|coreml|nnapi|cuda|xnnpack)"
        ),
    })
}

fn main() -> Result<()> {
    let args = Args::parse();
    let profile = parse_profile(&args.profile)?;
    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(parse_execution_provider)
        .transpose()?;
    let resolved_ep = match args.pipeline.as_str() {
        "v2" => explicit_ep.unwrap_or_else(polyvoice::onnx::ExecutionProvider::auto),
        _ => explicit_ep.unwrap_or(polyvoice::onnx::ExecutionProvider::Cpu),
    };

    // Build the requested pipeline. Each path verifies the integrity of exactly
    // the models it loads and returns the segmenter id for the report.
    let (mut runner, segmenter_id): (Runner, String) = match args.pipeline.as_str() {
        "v2" => {
            let clusterer = match args.clusterer.as_str() {
                "ahc" => ClustererKind::Ahc {
                    threshold: args.threshold,
                },
                "vbx" => ClustererKind::Vbx,
                other => anyhow::bail!("unknown --clusterer '{other}' (expected 'ahc' or 'vbx')"),
            };
            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,
                ..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 = V2Pipeline::builder()
                .config(cfg)
                .with_models_from(registry.clone())
                .build()
                .context("build v2 pipeline")?;
            (Runner::V2(Box::new(pipeline)), seg_id)
        }
        other => {
            if other != "legacy" {
                anyhow::bail!("unknown --pipeline '{other}' (expected 'legacy' or 'v2')");
            }
            let embedding_dim = profile.embedding_dim();
            let extractor =
                FbankOnnxExtractor::new(&models.embedder_path, embedding_dim, 1, resolved_ep)
                    .context("load embedder")?;
            let vad_path = registry.ensure("silero_vad").context("silero_vad model")?;
            let vad = SileroVad::new(&vad_path, 512).context("load vad")?;

            // Integrity gate: a DER number is only trustworthy if produced by
            // the EXACT shipped artifact — hard-fail if the on-disk embedder/VAD
            // sha256 disagrees with the manifest (swapped/corrupted/non-FP32).
            verify_model_integrity(&registry, profile, &models.embedder_path, &vad_path)?;

            let config = DiarizationConfig {
                cluster: ClusterConfig {
                    threshold: args.threshold,
                    min_cluster_size: args.min_cluster_size.unwrap_or(1),
                    min_cluster_secs: args.min_cluster_secs.unwrap_or(0.0),
                    ..Default::default()
                },
                ..DiarizationConfig::default()
            };
            let pipeline = Pipeline::new(config, VadConfig::default());
            (
                Runner::Legacy(Box::new(LegacyRunner {
                    pipeline,
                    extractor,
                    vad,
                })),
                "silero_vad".to_owned(),
            )
        }
    };

    // 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 audio_dir = args.dataset.join("audio");
    let rttm_dir = args.dataset.join("rttm");
    let mut wavs: Vec<PathBuf> = std::fs::read_dir(&audio_dir)
        .with_context(|| format!("read_dir {}", audio_dir.display()))?
        .filter_map(|e| e.ok())
        .filter(|e| e.path().extension().is_some_and(|x| x == "wav"))
        .map(|e| e.path())
        .collect();
    wavs.sort();
    if let Some(n) = args.max_files {
        wavs.truncate(n);
    }

    let dataset_name = args
        .dataset
        .file_name()
        .and_then(|s| s.to_str())
        .unwrap_or("unknown")
        .to_owned();

    let mut totals = Aggregate::default();
    let mut total_audio_secs = 0.0_f64;
    let mut total_runtime_secs = 0.0_f64;
    let mut stage_totals: Option<StageTimings> = None;
    // Per-file (collar, no-collar) DER pairs — the four report aggregates are
    // computed from these by the unit-tested aggregate_der helper.
    let mut der_pairs: Vec<(DerResult, DerResult)> = Vec::new();
    let mut speaker_count_exact = 0_usize;
    let mut speaker_count_pm1 = 0_usize;
    let mut speaker_count_off = 0_usize;
    let mut files_skipped = 0_usize;
    let mut per_file = Vec::with_capacity(wavs.len());

    for wav in &wavs {
        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");
            files_skipped += 1;
            continue;
        }
        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) = runner.run(&samples, sr)?;
        let runtime_secs = t0.elapsed().as_secs_f64();

        let ref_turns = {
            let raw = parse_rttm_file(&rttm).context("parse rttm")?;
            let grouped = group_by_file(&raw);
            // AMI-style fallback: EN2002a.Mix-Headset.wav → EN2002a
            let segs: Vec<_> = grouped
                .get(stem)
                .or_else(|| stem.split('.').next().and_then(|s| grouped.get(s)))
                .map(|v| v.iter().map(|s| (*s).clone()).collect())
                .unwrap_or_default();
            let (turns, _map) = to_speaker_turns(&segs);
            turns
        };

        // Headline collar + no-collar DER, restricted to the UEM scope when present
        // (AMI-style id fallback like the RTTM lookup). The overlap decomposition
        // below is a diagnostic and stays over the full file.
        let scored: Option<&[TimeRange]> = uem_map.as_ref().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) = 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 diff = ref_count.abs_diff(hyp_count);
        match diff {
            0 => speaker_count_exact += 1,
            1 => speaker_count_pm1 += 1,
            _ => speaker_count_off += 1,
        }

        totals.miss += der.miss_rate;
        totals.false_alarm += der.false_alarm_rate;
        totals.confusion += der.confusion_rate;
        totals.count += 1;
        der_pairs.push((der, der_no_collar));
        total_audio_secs += audio_secs;
        total_runtime_secs += runtime_secs;

        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(),
        );

        per_file.push(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,
        });
        if let Some(t) = stage_timings {
            let acc = 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;
        }
    }

    let n = totals.count.max(1) as f64;
    let agg = aggregate_der(&der_pairs);
    let der_collar_macro = agg.collar_macro;
    let der_no_collar_macro = agg.no_collar_macro;
    let der_collar_micro = agg.collar_micro;
    let der_no_collar_micro = agg.no_collar_micro;

    println!(
        "\n=== Aggregate DER over {} files (collar={:.2}s) ===",
        totals.count, args.collar
    );
    println!("  der_collar    : macro={der_collar_macro:.2}%  micro={der_collar_micro:.2}%");
    println!("  der_no_collar : macro={der_no_collar_macro:.2}%  micro={der_no_collar_micro:.2}%");

    let report = 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: totals.count,
        files_skipped,
        der_collar_macro,
        der_no_collar_macro,
        der_collar_micro,
        der_no_collar_micro,
        collar_secs: args.collar,
        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,
        miss: (totals.miss / n) * 100.0,
        false_alarm: (totals.false_alarm / n) * 100.0,
        confusion: (totals.confusion / n) * 100.0,
        rt_factor_avg: total_audio_secs / total_runtime_secs.max(1e-6),
        speaker_count: SpeakerCountDiagnostics {
            exact: speaker_count_exact,
            plus_minus_1: speaker_count_pm1,
            off_by_2_or_more: speaker_count_off,
        },
        model_hashes: model_hashes(&registry, profile, &segmenter_id),
        per_file,
    };
    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)]
mod prop_tests {
    use super::*;
    use proptest::prelude::*;

    /// Synthetic DerResult: `errors` error frames over `ref_frames` (all miss).
    fn synth(errors: u64, ref_frames: u64) -> DerResult {
        DerResult {
            der: errors as f64 / ref_frames as f64,
            miss_rate: errors as f64 / ref_frames as f64,
            false_alarm_rate: 0.0,
            confusion_rate: 0.0,
            total_speech: ref_frames as f64 * 0.01,
            total_ref_frames: ref_frames,
            missed_frames: errors,
            false_alarm_frames: 0,
            confusion_frames: 0,
        }
    }

    #[test]
    fn aggregate_macro_diverges_from_micro_and_micro_is_frame_weighted() {
        // A tiny 1s file at 50% DER and a long 60s file at 1% DER: the mean of
        // ratios (macro) must NOT equal the frame-weighted micro, and micro
        // must equal summed error frames / summed reference frames exactly.
        let short = synth(50, 100);
        let long = synth(60, 6000);
        let agg = aggregate_der(&[(short, short), (long, long)]);
        assert!(
            (agg.collar_macro - 25.5).abs() < 1e-9,
            "{}",
            agg.collar_macro
        );
        let expected_micro = (50 + 60) as f64 / (100 + 6000) as f64 * 100.0;
        assert!((agg.collar_micro - expected_micro).abs() < 1e-9);
        assert!((agg.collar_macro - agg.collar_micro).abs() > 10.0);
        // Same inputs on both passes => identical aggregates per pass.
        assert_eq!(agg.collar_micro, agg.no_collar_micro);
    }

    #[test]
    fn aggregate_no_collar_at_least_collar_on_boundary_errors() {
        use polyvoice::types::{SpeakerId, SpeakerTurn, TimeRange};
        let turn = |s: u32, a: f64, b: f64| SpeakerTurn {
            speaker: SpeakerId(s),
            time: TimeRange { start: a, end: b },
            text: None,
        };
        // Hypothesis shifted 0.3s off every reference boundary: the collar
        // forgives part of that error, no-collar must not.
        let reference = vec![turn(0, 0.0, 10.0), turn(1, 12.0, 20.0)];
        let hypothesis = vec![turn(0, 0.3, 10.3), turn(1, 12.3, 20.3)];
        let collar = compute_der(&reference, &hypothesis, 0.25);
        let no_collar = compute_der(&reference, &hypothesis, 0.0);
        let agg = aggregate_der(&[(collar, no_collar)]);
        assert!(
            agg.no_collar_micro >= agg.collar_micro,
            "no-collar {} < collar {}",
            agg.no_collar_micro,
            agg.collar_micro
        );
        assert!(agg.no_collar_macro >= agg.collar_macro);
        assert!(agg.no_collar_micro > 0.0, "boundary errors must be scored");
    }

    proptest! {
        #[test]
        fn bench_args_parses_with_valid_args(
            profile in "(mobile|balanced)",
            collar in 0.0f64..1.0f64,
            threshold in 0.0f32..1.0f32,
            max_files in 0usize..100usize,
        ) {
            let args = vec![
                "polyvoice-bench".to_string(),
                "/tmp/dataset".to_string(),
                "--profile".to_string(), profile,
                "--collar".to_string(), collar.to_string(),
                "--threshold".to_string(), threshold.to_string(),
                "--max-files".to_string(), max_files.to_string(),
            ];
            let result = Args::try_parse_from(&args);
            prop_assert!(result.is_ok());
        }

        #[test]
        fn parse_profile_accepts_only_valid(s in "[a-zA-Z0-9_-]{1,20}") {
            let result = parse_profile(&s);
            if s == "mobile" || s == "balanced" {
                prop_assert!(result.is_ok());
            } else {
                prop_assert!(result.is_err());
            }
        }
    }
}