polyvoice 0.8.0

//! Sliding-window aggregator for powerset segmentation outputs.
//!
//! Combines per-window 7-class logits into file-globally-consistent
//! `RawSegment` outputs.
//!
//! Algorithm:
//! 1. For each adjacent window pair (i, i+1), build the 3×3 IoU matrix between
//!    speaker masks in the temporal overlap region. Each speaker mask is the
//!    union of frames where the window's argmax label includes that speaker.
//! 2. Use Kuhn-Munkres on `-IoU` to find the assignment that maps window i+1's
//!    local indices onto window i's. Apply the permutation so the same person
//!    has the same index file-wide.
//! 3. For every audio frame, average the per-class logits across every window
//!    that contains that frame.
//! 4. Argmax each averaged logit vector → frame label.
//! 5. Run-length encode consecutive identical labels into `RawSegment`s.

use crate::hungarian;
use crate::segmentation::decoder::{FrameLabel, PowersetClass, PowersetDecoder};
use crate::segmentation::{RawSegment, SegmentationError};
use crate::types::TimeRange;

/// One window's segmentation output.
#[derive(Debug, Clone)]
pub struct WindowOutput {
    pub start_time: f32,
    pub end_time: f32,
    /// Row-major `(num_frames, 7)` logits.
    pub logits: Vec<f32>,
    pub num_frames: usize,
}

impl WindowOutput {
    /// { true }
    /// `pub fn new( start_time: f32, end_time: f32, logits: Vec<f32>, num_frames: usize, ) -> Result<Self, SegmentationError>`
    /// { ret.as_ref().map_or(true, |w| w.logits.len() == w.num_frames * 7) }
    pub fn new(
        start_time: f32,
        end_time: f32,
        logits: Vec<f32>,
        num_frames: usize,
    ) -> Result<Self, SegmentationError> {
        if logits.len() != num_frames * 7 {
            return Err(SegmentationError::InvalidOutputShape {
                actual_shape: vec![logits.len()],
            });
        }
        Ok(Self {
            start_time,
            end_time,
            logits,
            num_frames,
        })
    }

    /// { true }
    /// pub fn frame_stride(&self) -> f32
    /// { ret >= 0.0 || self.num_frames == 0 }
    pub fn frame_stride(&self) -> f32 {
        if self.num_frames == 0 {
            0.0
        } else {
            (self.end_time - self.start_time) / self.num_frames as f32
        }
    }

    /// { true }
    /// pub fn frame_time(&self, frame_idx: usize) -> f32
    /// { ret == self.start_time + frame_idx as f32 * self.frame_stride() }
    pub fn frame_time(&self, frame_idx: usize) -> f32 {
        self.start_time + frame_idx as f32 * self.frame_stride()
    }
}

/// Configuration for aggregation.
#[derive(Debug, Clone)]
pub struct AggregationConfig {
    pub min_segment_secs: f32,
    pub max_local_speakers: usize,
}

impl Default for AggregationConfig {
    fn default() -> Self {
        Self {
            min_segment_secs: 0.0,
            max_local_speakers: 3,
        }
    }
}

/// Aggregator over sliding-window powerset outputs.
pub struct Aggregator {
    config: AggregationConfig,
}

impl Aggregator {
    /// { true }
    /// pub fn new(config: AggregationConfig) -> Self
    /// { true }
    pub fn new(config: AggregationConfig) -> Self {
        Self { config }
    }

    /// { true }
    /// pub fn config(&self) -> &AggregationConfig
    /// { ret == &self.config }
    pub fn config(&self) -> &AggregationConfig {
        &self.config
    }

    /// { true }
    /// `pub fn stitch(&self, windows: &[WindowOutput]) -> Result<Vec<RawSegment>, SegmentationError>`
    /// { true }
    /// Stitch overlapping windows into file-consistent `RawSegment`s.
    pub fn stitch(&self, windows: &[WindowOutput]) -> Result<Vec<RawSegment>, SegmentationError> {
        if windows.is_empty() {
            return Ok(Vec::new());
        }

        // 1) For each window, compute per-frame argmax labels.
        let mut window_labels: Vec<Vec<FrameLabel>> = Vec::with_capacity(windows.len());
        for w in windows {
            let labels = PowersetDecoder::decode_window(&w.logits, w.num_frames)?;
            window_labels.push(labels);
        }

        // 2) Hungarian-align each adjacent window pair: permute window i+1's local
        // speaker indices onto window i's reference frame.
        let mut permutations: Vec<[u8; 3]> =
            std::iter::repeat_n([0u8, 1u8, 2u8], windows.len()).collect();

        for i in 1..windows.len() {
            let perm = self.window_permutation(
                &windows[i - 1],
                &window_labels[i - 1],
                &windows[i],
                &window_labels[i],
                &permutations[i - 1],
            )?;
            // `window_permutation` already built A's masks using
            // `a_perm_so_far`, so `perm` maps window i's local indices directly
            // onto the file-global speaker frame. No extra composition needed.
            permutations[i] = perm;
        }

        // 3-5) For every audio frame across the file, average per-class logits over
        // every window that contains it. Argmax + run-length encode.
        self.average_and_run_length_encode(windows, &window_labels, &permutations)
    }

    /// Compute the permutation that maps window B's local indices onto A's frame.
    fn window_permutation(
        &self,
        a: &WindowOutput,
        a_labels: &[FrameLabel],
        b: &WindowOutput,
        b_labels: &[FrameLabel],
        a_perm_so_far: &[u8; 3],
    ) -> Result<[u8; 3], SegmentationError> {
        let n = self.config.max_local_speakers.min(3);
        let overlap_start = a.start_time.max(b.start_time);
        let overlap_end = a.end_time.min(b.end_time);
        if overlap_end <= overlap_start || n == 0 {
            return Ok([0, 1, 2]);
        }

        let stride = a.frame_stride().max(1e-6);
        let grid_len = ((overlap_end - overlap_start) / stride).ceil() as usize;
        if grid_len == 0 {
            return Ok([0, 1, 2]);
        }

        let mut a_masks = vec![vec![false; grid_len]; 3];
        let mut b_masks = vec![vec![false; grid_len]; 3];
        for k in 0..grid_len {
            let t = overlap_start + k as f32 * stride;
            if let Some(idx_a) = self.frame_index_at(a, t) {
                if idx_a < a_labels.len() {
                    for s in a_labels[idx_a].class.speakers() {
                        if (s as usize) < 3 {
                            let permuted = a_perm_so_far[s as usize] as usize;
                            if permuted < 3 {
                                a_masks[permuted][k] = true;
                            }
                        }
                    }
                }
            }
            if let Some(idx_b) = self.frame_index_at(b, t) {
                if idx_b < b_labels.len() {
                    for s in b_labels[idx_b].class.speakers() {
                        if (s as usize) < 3 {
                            b_masks[s as usize][k] = true;
                        }
                    }
                }
            }
        }

        // Build cost matrix: C[a][b] = -IoU(a_mask, b_mask).
        let mut cost: Vec<Vec<f32>> = vec![vec![0.0_f32; n]; n];
        let a_active_count = a_masks
            .iter()
            .take(n)
            .filter(|m| m.iter().any(|&x| x))
            .count();
        let b_active_count = b_masks
            .iter()
            .take(n)
            .filter(|m| m.iter().any(|&x| x))
            .count();

        for ai in 0..n {
            for bi in 0..n {
                let mut inter = 0_usize;
                let mut uni = 0_usize;
                for k in 0..grid_len {
                    let ax = a_masks[ai][k];
                    let bx = b_masks[bi][k];
                    if ax && bx {
                        inter += 1;
                    }
                    if ax || bx {
                        uni += 1;
                    }
                }
                let iou = if uni == 0 {
                    0.0
                } else {
                    inter as f32 / uni as f32
                };
                cost[ai][bi] = -iou;
            }
        }

        // If fewer than 2 speakers are active on either side in the overlap, we
        // cannot reliably determine the full permutation — return identity.
        if a_active_count < 2 || b_active_count < 2 {
            return Ok([0, 1, 2]);
        }

        let assignment =
            hungarian::solve(&cost).ok_or_else(|| SegmentationError::PermutationFailed {
                prev_idx: 0,
                next_idx: 0,
                detail: "non-square cost matrix".to_owned(),
            })?;

        // assignment[i] = j means: row i (A's global speaker i) best matches column j
        // (B's local speaker j). We want perm[b_local] = a_global, i.e. the direct
        // mapping: perm[j] = i.
        let mut perm = [0_u8, 1_u8, 2_u8];
        for (i, &j) in assignment.iter().enumerate() {
            if j < 3 && i < 3 {
                perm[j] = i as u8;
            }
        }
        Ok(perm)
    }

    /// Find the frame index in `w` whose center is closest to time `t`. Returns
    /// `None` if `t` is outside the window's span.
    ///
    /// Uses `floor((t - start)/stride)`. This already IS the nearest-center frame:
    /// `round((t - start)/stride - 0.5)` equals this `floor` for every in-span `t`
    /// once clamped to `[0, num_frames-1]`, so it matches the center convention used
    /// by `average_and_run_length_encode` (proven by
    /// `frame_index_floor_equals_nearest_center`). Do NOT "fix" this to `round`; it
    /// is a no-op that only changes the out-of-span `t == start - ε` corner.
    fn frame_index_at(&self, w: &WindowOutput, t: f32) -> Option<usize> {
        if t < w.start_time || t > w.end_time || w.num_frames == 0 {
            return None;
        }
        let stride = w.frame_stride();
        if stride <= 0.0 {
            return None;
        }
        let idx = ((t - w.start_time) / stride).floor() as usize;
        Some(idx.min(w.num_frames - 1))
    }

    /// Average per-class logits across windows that contain each global frame,
    /// then argmax + run-length encode into `RawSegment`s.
    fn average_and_run_length_encode(
        &self,
        windows: &[WindowOutput],
        window_labels: &[Vec<FrameLabel>],
        permutations: &[[u8; 3]],
    ) -> Result<Vec<RawSegment>, SegmentationError> {
        let stride = windows[0].frame_stride().max(1e-6);
        let global_start = windows
            .iter()
            .map(|w| w.start_time)
            .fold(f32::INFINITY, f32::min);
        let global_end = windows
            .iter()
            .map(|w| w.end_time)
            .fold(f32::NEG_INFINITY, f32::max);
        let global_frames = ((global_end - global_start) / stride).ceil() as usize;

        let mut summed_probs = vec![[0.0_f32; 7]; global_frames];
        let mut counts = vec![0_u32; global_frames];

        for (wi, w) in windows.iter().enumerate() {
            let perm = permutations[wi];
            for f in 0..w.num_frames {
                let t_center = w.frame_time(f) + 0.5 * stride;
                let g_idx_f = (t_center - global_start) / stride;
                if g_idx_f < 0.0 {
                    continue;
                }
                let g_idx = g_idx_f.floor() as usize;
                if g_idx >= global_frames {
                    continue;
                }
                if window_labels[wi].get(f).is_none() {
                    continue;
                }

                let frame_logits = &w.logits[f * 7..(f + 1) * 7];

                let mut max_logit = f32::NEG_INFINITY;
                for &l in frame_logits {
                    if l > max_logit {
                        max_logit = l;
                    }
                }
                let mut exps = [0.0_f32; 7];
                let mut sum = 0.0_f32;
                for (i, &l) in frame_logits.iter().enumerate() {
                    exps[i] = (l - max_logit).exp();
                    sum += exps[i];
                }
                let inv_sum = if sum > 0.0 { 1.0 / sum } else { 1.0 };

                // Apply permutation: rebuild the softmax vector under the file-global
                // speaker ordering by remapping each class's speaker set through `perm`.
                let mut remapped = [0.0_f32; 7];
                for (c, _) in exps.iter().enumerate() {
                    if let Some(class) = PowersetDecoder::class_for_index(c) {
                        let speakers = class.speakers();
                        let remapped_speakers: Vec<u8> = speakers
                            .iter()
                            .map(|s| {
                                if (*s as usize) < 3 {
                                    perm[*s as usize]
                                } else {
                                    *s
                                }
                            })
                            .collect();
                        let new_class = match remapped_speakers.as_slice() {
                            [] => 0,
                            [s] => 1 + (*s as usize),
                            [a, b] => {
                                let (lo, hi) = if a < b {
                                    (*a as usize, *b as usize)
                                } else {
                                    (*b as usize, *a as usize)
                                };
                                match (lo, hi) {
                                    (0, 1) => 4,
                                    (0, 2) => 5,
                                    (1, 2) => 6,
                                    _ => 0,
                                }
                            }
                            _ => 0,
                        };
                        remapped[new_class] += exps[c] * inv_sum;
                    }
                }

                for (i, &p) in remapped.iter().enumerate() {
                    summed_probs[g_idx][i] += p;
                }
                counts[g_idx] += 1;
            }
        }

        let mut frame_classes: Vec<Option<PowersetClass>> = Vec::with_capacity(global_frames);
        let mut frame_confidences: Vec<f32> = Vec::with_capacity(global_frames);
        for g in 0..global_frames {
            if counts[g] == 0 {
                frame_classes.push(None);
                frame_confidences.push(0.0);
                continue;
            }
            let inv = 1.0 / counts[g] as f32;
            let mut argmax = 0_usize;
            let mut maxp = 0.0_f32;
            for (c, &sp) in summed_probs[g].iter().enumerate() {
                let p = sp * inv;
                if p > maxp {
                    maxp = p;
                    argmax = c;
                }
            }
            frame_classes.push(PowersetDecoder::class_for_index(argmax));
            frame_confidences.push(maxp);
        }

        // Run-length encode per speaker. A run is broken not only when the
        // speaker falls silent but also when its overlap status flips: a speaker
        // talking across a brief overlap emits separate solo and overlap
        // segments rather than one run tainted by the overlap. This keeps
        // `is_overlap` precise (a segment is overlap iff *every* frame in it was
        // an overlap frame) and makes two simultaneously-active speakers emit
        // time-equal overlap segments that `extract_overlap_time_ranges` pairs.
        let mut segments: Vec<RawSegment> = Vec::new();
        // (start_global_frame, confidence_sum, confidence_count, is_overlap_run)
        let mut active: [Option<(usize, f32, f32, bool)>; 3] = [None, None, None];

        for g in 0..global_frames {
            let frame_class = frame_classes[g];
            let conf = frame_confidences[g];
            let is_overlap_frame = frame_class.map(|c| c.is_overlap()).unwrap_or(false);
            let active_speakers: Vec<u8> = match frame_class {
                Some(c) => c.speakers(),
                None => Vec::new(),
            };

            for (s, slot) in active.iter_mut().enumerate() {
                let s_active_now = active_speakers.iter().any(|x| *x as usize == s);
                let ov_now = s_active_now && is_overlap_frame;
                match (*slot, s_active_now) {
                    (None, true) => {
                        *slot = Some((g, conf, 1.0, ov_now));
                    }
                    (Some((start_g, conf_sum, conf_count, run_ov)), true) if run_ov == ov_now => {
                        *slot = Some((start_g, conf_sum + conf, conf_count + 1.0, run_ov));
                    }
                    (Some((start_g, conf_sum, conf_count, run_ov)), true) => {
                        // Overlap status flipped — close the current run, open a new one at g.
                        push_segment(
                            &mut segments,
                            global_start,
                            stride,
                            s,
                            start_g,
                            g,
                            conf_sum,
                            conf_count,
                            run_ov,
                            self.config.min_segment_secs,
                        );
                        *slot = Some((g, conf, 1.0, ov_now));
                    }
                    (Some((start_g, conf_sum, conf_count, run_ov)), false) => {
                        push_segment(
                            &mut segments,
                            global_start,
                            stride,
                            s,
                            start_g,
                            g,
                            conf_sum,
                            conf_count,
                            run_ov,
                            self.config.min_segment_secs,
                        );
                        *slot = None;
                    }
                    (None, false) => {}
                }
            }
        }

        // Flush trailing active runs.
        for (s, slot) in active.iter().enumerate() {
            if let Some((start_g, conf_sum, conf_count, run_ov)) = *slot {
                push_segment(
                    &mut segments,
                    global_start,
                    stride,
                    s,
                    start_g,
                    global_frames,
                    conf_sum,
                    conf_count,
                    run_ov,
                    self.config.min_segment_secs,
                );
            }
        }

        segments.sort_by(|a, b| {
            a.time
                .start
                .partial_cmp(&b.time.start)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        Ok(segments)
    }
}

/// Emit one run as a `RawSegment`, skipping runs shorter than `min_segment_secs`.
/// `start_g`/`end_g` are global frame indices; the run carries a single
/// `is_overlap` flag because the RLE splits runs at every overlap-status change.
#[allow(clippy::too_many_arguments)]
fn push_segment(
    segments: &mut Vec<RawSegment>,
    global_start: f32,
    stride: f32,
    speaker: usize,
    start_g: usize,
    end_g: usize,
    conf_sum: f32,
    conf_count: f32,
    is_overlap: bool,
    min_segment_secs: f32,
) {
    let start_t = global_start + start_g as f32 * stride;
    let end_t = global_start + end_g as f32 * stride;
    if end_t - start_t < min_segment_secs {
        return;
    }
    let mean_conf = (conf_sum / conf_count.max(1.0)).clamp(0.0, 1.0);
    segments.push(RawSegment {
        time: TimeRange {
            start: start_t as f64,
            end: end_t as f64,
        },
        local_speaker_idx: speaker as u8,
        is_overlap,
        confidence: PowersetDecoder::frame_confidence(mean_conf),
    });
}

#[allow(clippy::unwrap_used)]
#[cfg(test)]
mod tests {
    use super::*;

    /// Helper: build a window where every frame is a single class (like 0=silence,
    /// 1=speaker 0, etc.) with the listed class as logit 10 and others as logit 0.
    fn synthetic_window(
        start: f32,
        end: f32,
        num_frames: usize,
        classes: &[usize],
    ) -> WindowOutput {
        assert_eq!(classes.len(), num_frames);
        let mut logits = Vec::with_capacity(num_frames * 7);
        for &c in classes {
            for k in 0..7 {
                logits.push(if k == c { 10.0 } else { 0.0 });
            }
        }
        WindowOutput::new(start, end, logits, num_frames).unwrap()
    }

    /// Frame-time convention check: `frame_index_at` uses
    /// `floor((t - start)/stride)`, and the RLE pass (line ~300) places frame `f`
    /// by its center `start + (f + 0.5)*stride`. These are NOT two different
    /// conventions: `floor(x)` already returns the frame whose CENTER is closest to
    /// `t`, because `round(x - 0.5) == floor(x)` for every non-negative `x` once the
    /// result is clamped to `[0, num_frames-1]`. This test pins that equivalence so
    /// a future "fix" to `round((t-start)/stride - 0.5)` is recognized as a no-op.
    #[test]
    fn frame_index_floor_equals_nearest_center() {
        let stride = 0.1f32;
        let start = 0.37f32;
        for i in 0..5000 {
            let t = start + i as f32 * 0.00713;
            let x = (t - start) / stride;
            // Clamp both at the lower edge the way frame_index_at does.
            let floor_idx = (x.floor() as i64).max(0);
            let round_idx = ((x - 0.5).round() as i64).max(0);
            assert_eq!(
                floor_idx, round_idx,
                "floor and nearest-center disagree at t={t} x={x:.6}"
            );
        }
    }

    /// Frame-time convention check: a speaker change staggered ~0.5*stride off a
    /// window boundary must still be labelled consistently after stitching. With the
    /// sampler (`frame_index_at`) and the RLE applier sharing the nearest-center
    /// convention, there is no 1-frame boundary flip — this passes on current code,
    /// confirming the two conventions already coincide (no off-by-one).
    #[test]
    fn staggered_speaker_change_is_labelled_consistently() {
        // Window A: 0.0–5.0, 50 frames (stride 0.1). spk0 (class 1) then spk1
        // (class 2); change at frame 25 → t = 2.5s.
        let mut a_classes = vec![1usize; 50];
        for c in &mut a_classes[25..50] {
            *c = 2;
        }
        let a = synthetic_window(0.0, 5.0, 50, &a_classes);
        // Window B: 2.45–7.45, 50 frames (stride 0.1) — its grid is offset half a
        // stride from A. Same physical truth: spk0 until ~2.5s then spk1.
        let mut b_classes = vec![1usize; 50];
        for c in &mut b_classes[1..50] {
            *c = 2;
        }
        let b = synthetic_window(2.45, 7.45, 50, &b_classes);

        let agg = Aggregator::new(AggregationConfig::default());
        let segs = agg.stitch(&[a, b]).unwrap();

        fn speaker_at(segs: &[RawSegment], t: f64) -> Option<u8> {
            segs.iter()
                .find(|s| s.time.start <= t && t < s.time.end)
                .map(|s| s.local_speaker_idx)
        }
        // Well clear of the staggered boundary, the two speakers are distinct and
        // stable (identity permutation — both windows use the same class indices).
        let early = speaker_at(&segs, 1.0).expect("segment around 1.0s");
        let late = speaker_at(&segs, 6.0).expect("segment around 6.0s");
        assert_ne!(
            early, late,
            "the two speakers must stay distinct across the seam"
        );
        // Exactly two global speakers across the file.
        let unique: std::collections::HashSet<u8> =
            segs.iter().map(|s| s.local_speaker_idx).collect();
        assert_eq!(unique.len(), 2, "expected 2 speakers, got {}", unique.len());
    }

    #[test]
    fn empty_returns_empty() {
        let agg = Aggregator::new(AggregationConfig::default());
        assert!(agg.stitch(&[]).unwrap().is_empty());
    }

    #[test]
    fn single_window_silence_yields_no_segments() {
        let agg = Aggregator::new(AggregationConfig::default());
        let w = synthetic_window(0.0, 1.0, 10, &[0; 10]);
        let segs = agg.stitch(&[w]).unwrap();
        assert!(segs.is_empty());
    }

    #[test]
    fn single_window_one_speaker_yields_one_segment() {
        let agg = Aggregator::new(AggregationConfig::default());
        let w = synthetic_window(0.0, 1.0, 10, &[1; 10]);
        let segs = agg.stitch(&[w]).unwrap();
        assert_eq!(segs.len(), 1);
        assert_eq!(segs[0].local_speaker_idx, 0);
        assert!(!segs[0].is_overlap);
    }

    #[test]
    fn single_window_overlap_yields_two_segments_same_time() {
        let agg = Aggregator::new(AggregationConfig::default());
        let w = synthetic_window(0.0, 1.0, 10, &[4; 10]);
        let segs = agg.stitch(&[w]).unwrap();
        assert_eq!(segs.len(), 2);
        assert!((segs[0].time.start - segs[1].time.start).abs() < 1e-3);
        assert!((segs[0].time.end - segs[1].time.end).abs() < 1e-3);
        assert!(segs.iter().all(|s| s.is_overlap));
        let speakers: Vec<u8> = segs.iter().map(|s| s.local_speaker_idx).collect();
        assert!(speakers.contains(&0));
        assert!(speakers.contains(&1));
    }

    #[test]
    fn partial_overlap_run_splits_into_solo_and_overlap_segments() {
        // spk0 talks the whole 0-10s window; spk1 joins only over 4-6s (class 4 =
        // pair{0,1}). spk0's run must split into solo [0,4), overlap [4,6), solo
        // [6,10); spk1 emits one overlap [4,6). The two overlap pieces must share
        // an exact time range so extract_overlap_time_ranges can pair them — this
        // is the fix for whole single-speaker runs being falsely flagged overlap.
        let mut classes = vec![1usize; 100];
        for c in &mut classes[40..60] {
            *c = 4;
        }
        let w = synthetic_window(0.0, 10.0, 100, &classes);
        let agg = Aggregator::new(AggregationConfig::default());
        let segs = agg.stitch(&[w]).unwrap();

        let overlap: Vec<&RawSegment> = segs.iter().filter(|s| s.is_overlap).collect();
        assert_eq!(
            overlap.len(),
            2,
            "exactly two overlap pieces (one per speaker)"
        );
        for s in &overlap {
            assert!((s.time.start - 4.0).abs() < 1e-3, "overlap starts at 4.0s");
            assert!((s.time.end - 6.0).abs() < 1e-3, "overlap ends at 6.0s");
        }
        let ov_speakers: std::collections::HashSet<u8> =
            overlap.iter().map(|s| s.local_speaker_idx).collect();
        assert_eq!(ov_speakers, [0u8, 1u8].into_iter().collect());

        // spk0: three pieces (solo, overlap, solo); the two solo ones are NOT overlap.
        let spk0: Vec<&RawSegment> = segs.iter().filter(|s| s.local_speaker_idx == 0).collect();
        assert_eq!(spk0.len(), 3, "spk0 run splits at both overlap boundaries");
        let solo0: Vec<&&RawSegment> = spk0.iter().filter(|s| !s.is_overlap).collect();
        assert_eq!(solo0.len(), 2, "spk0 keeps two solo pieces");
    }

    #[test]
    fn two_windows_with_consistent_speakers_remain_consistent() {
        let a = synthetic_window(0.0, 5.0, 50, &[1; 50]);
        let b = synthetic_window(4.0, 9.0, 50, &[1; 50]);
        let agg = Aggregator::new(AggregationConfig::default());
        let segs = agg.stitch(&[a, b]).unwrap();
        assert!(segs.iter().all(|s| s.local_speaker_idx == 0));
        assert!(segs.iter().all(|s| !s.is_overlap));
    }

    #[test]
    fn two_windows_requiring_permutation_get_aligned() {
        let a = synthetic_window(
            0.0,
            5.0,
            50,
            &[
                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
            ],
        );
        let b = synthetic_window(
            4.0,
            9.0,
            50,
            &[
                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
                2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
            ],
        );
        let agg = Aggregator::new(AggregationConfig::default());
        let segs = agg.stitch(&[a, b]).unwrap();

        let mut idx_set = std::collections::HashSet::new();
        for s in &segs {
            idx_set.insert(s.local_speaker_idx);
        }
        assert_eq!(idx_set.len(), 2);

        let mut sorted = segs.clone();
        sorted.sort_by(|a, b| a.time.start.partial_cmp(&b.time.start).unwrap());
        let first = sorted.first().unwrap();
        let last = sorted.last().unwrap();
        assert_ne!(first.local_speaker_idx, last.local_speaker_idx);
    }

    /// Regression test for the cumulative-permutation double-application bug.
    /// Windows 0 and 1 are swapped once; window 2 is swapped once relative to
    /// window 1. Because `window_permutation` already applies the cumulative
    /// permutation when building A-masks, the returned `perm` is already
    /// file-global. Before the fix the code composed `prev[perm[...]]`, which
    /// double-applied the permutation and produced inconsistent global speaker
    /// indices across window boundaries.
    #[test]
    fn three_windows_keep_global_speaker_indices_consistent() {
        // Window 0: spk0 in 0-3.0s and 4.0-4.5s; spk1 in 3.0-4.0s and 4.5-5.0s.
        // The 4.0-5.0s overlap with window 1 contains both speakers.
        let w0 = synthetic_window(
            0.0,
            5.0,
            50,
            &[
                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2,
            ],
        );
        // Window 1 overlaps 4-5s. In the overlap it swaps local indices
        // relative to window 0, so local spk1 = global spk0 and local spk0 =
        // global spk1. Both speakers remain active through the window so the
        // overlap with window 2 also contains both global speakers.
        let w1 = synthetic_window(
            4.0,
            9.0,
            50,
            &[
                2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
                2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1,
            ],
        );
        // Window 2 overlaps 8-9s. In the overlap it swaps local indices
        // relative to window 1, so local spk1 = global spk0 and local spk0 =
        // global spk1. In the non-overlap region only global spk1 continues.
        let w2 = synthetic_window(
            8.0,
            13.0,
            50,
            &[
                2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
            ],
        );

        let agg = Aggregator::new(AggregationConfig::default());
        let segs = agg.stitch(&[w0, w1, w2]).unwrap();

        fn speaker_at(segs: &[RawSegment], t: f64) -> Option<u8> {
            segs.iter()
                .find(|s| s.time.start <= t && t < s.time.end)
                .map(|s| s.local_speaker_idx)
        }

        let spk_0_early = speaker_at(&segs, 1.0).expect("segment expected around 1.0s");
        let spk_0_late = speaker_at(&segs, 8.25).expect("segment expected around 8.25s");
        let spk_1_early = speaker_at(&segs, 7.5).expect("segment expected around 7.5s");
        let spk_1_late = speaker_at(&segs, 11.0).expect("segment expected around 11.0s");

        assert_eq!(
            spk_0_early, spk_0_late,
            "global speaker 0 must keep the same index across window 2 boundary"
        );
        assert_eq!(
            spk_1_early, spk_1_late,
            "global speaker 1 must keep the same index across window 2 boundary"
        );
        assert_ne!(spk_0_early, spk_1_early, "two distinct speakers expected");
    }

    #[test]
    fn min_segment_filter_drops_tiny_runs() {
        let w = synthetic_window(0.0, 1.0, 100, &{
            let mut v = vec![0; 100];
            v[50] = 1;
            v
        });
        let config = AggregationConfig {
            min_segment_secs: 0.1,
            ..AggregationConfig::default()
        };
        let agg = Aggregator::new(config);
        let segs = agg.stitch(&[w]).unwrap();
        assert!(segs.is_empty());
    }

    #[test]
    fn output_segments_are_sorted_by_start_time() {
        let mut classes = vec![0; 100];
        for c in &mut classes[10..20] {
            *c = 1;
        }
        for c in &mut classes[50..60] {
            *c = 1;
        }
        let w = synthetic_window(0.0, 1.0, 100, &classes);
        let agg = Aggregator::new(AggregationConfig::default());
        let segs = agg.stitch(&[w]).unwrap();
        assert!(segs.len() >= 2);
        for pair in segs.windows(2) {
            assert!(pair[0].time.start <= pair[1].time.start);
        }
    }

    #[test]
    fn confidence_is_within_unit_interval() {
        let w = synthetic_window(0.0, 1.0, 10, &[1; 10]);
        let agg = Aggregator::new(AggregationConfig::default());
        let segs = agg.stitch(&[w]).unwrap();
        for s in segs {
            assert!(s.confidence.get() >= 0.0);
            assert!(s.confidence.get() <= 1.0);
        }
    }
}