//! Subcarrier partitioning via graph min-cut (ruvector-mincut).
//!
//! Uses [`MinCutBuilder`] to partition subcarriers into two groups —
//! **sensitive** (high body-motion correlation) and **insensitive** (dominated
//! by static multipath or noise) — based on pairwise sensitivity similarity.
//!
//! The edge weight between subcarriers `i` and `j` is the inverse absolute
//! difference of their sensitivity scores; highly similar subcarriers have a
//! heavy edge, making the min-cut prefer to separate dissimilar ones.
//!
//! A virtual source (node `n`) and sink (node `n+1`) are added to make the
//! graph connected and enable the min-cut to naturally bifurcate the
//! subcarrier set. The cut edges that cross from the source-side to the
//! sink-side identify the two partitions.

use ruvector_mincut::{DynamicMinCut, MinCutBuilder};

/// Partition `sensitivity` scores into (sensitive_indices, insensitive_indices)
/// using graph min-cut. The group with higher mean sensitivity is "sensitive".
///
/// # Arguments
///
/// - `sensitivity`: per-subcarrier sensitivity score, one value per subcarrier.
///   Higher values indicate stronger body-motion correlation.
///
/// # Returns
///
/// A tuple `(sensitive, insensitive)` where each element is a `Vec<usize>` of
/// subcarrier indices belonging to that partition. Together they cover all
/// indices `0..sensitivity.len()`.
///
/// # Notes
///
/// When `sensitivity` is empty or all edges would be below threshold the
/// function falls back to a simple midpoint split.
pub fn mincut_subcarrier_partition(sensitivity: &[f32]) -> (Vec<usize>, Vec<usize>) {
    let n = sensitivity.len();
    if n == 0 {
        return (Vec::new(), Vec::new());
    }
    if n == 1 {
        return (vec![0], Vec::new());
    }

    // Build edges as a flow network:
    // - Nodes 0..n-1 are subcarrier nodes
    // - Node n is the virtual source (connected to high-sensitivity nodes)
    // - Node n+1 is the virtual sink (connected to low-sensitivity nodes)
    let source = n as u64;
    let sink = (n + 1) as u64;

    let mean_sens: f32 = sensitivity.iter().sum::<f32>() / n as f32;

    let mut edges: Vec<(u64, u64, f64)> = Vec::new();

    // Source connects to subcarriers with above-average sensitivity.
    // Sink connects to subcarriers with below-average sensitivity.
    for (i, &sens) in sensitivity.iter().enumerate().take(n) {
        let cap = (sens as f64).abs() + 1e-6;
        if sens >= mean_sens {
            edges.push((source, i as u64, cap));
        } else {
            edges.push((i as u64, sink, cap));
        }
    }

    // Subcarrier-to-subcarrier edges weighted by inverse sensitivity difference.
    let threshold = 0.1_f64;
    for i in 0..n {
        for j in (i + 1)..n {
            let diff = (sensitivity[i] - sensitivity[j]).abs() as f64;
            let weight = if diff > 1e-9 { 1.0 / diff } else { 1e6_f64 };
            if weight > threshold {
                edges.push((i as u64, j as u64, weight));
                edges.push((j as u64, i as u64, weight));
            }
        }
    }

    let mc: DynamicMinCut = match MinCutBuilder::new().exact().with_edges(edges).build() {
        Ok(mc) => mc,
        Err(_) => {
            // Fallback: midpoint split on builder error.
            let mid = n / 2;
            return ((0..mid).collect(), (mid..n).collect());
        }
    };

    // Use cut_edges to identify which side each node belongs to.
    // Nodes reachable from source in the residual graph are "source-side",
    // the rest are "sink-side".
    let cut = mc.cut_edges();

    // Collect nodes that appear on the source side of a cut edge (u nodes).
    let mut source_side: std::collections::HashSet<u64> = std::collections::HashSet::new();
    let mut sink_side: std::collections::HashSet<u64> = std::collections::HashSet::new();

    for edge in &cut {
        // Cut edge goes from source-side node to sink-side node.
        if edge.source != source && edge.source != sink {
            source_side.insert(edge.source);
        }
        if edge.target != source && edge.target != sink {
            sink_side.insert(edge.target);
        }
    }

    // Any subcarrier not explicitly classified goes to whichever side is smaller.
    let mut side_a: Vec<usize> = source_side.iter().map(|&x| x as usize).collect();
    let mut side_b: Vec<usize> = sink_side.iter().map(|&x| x as usize).collect();

    // Assign unclassified nodes.
    for i in 0..n {
        if !source_side.contains(&(i as u64)) && !sink_side.contains(&(i as u64)) {
            if side_a.len() <= side_b.len() {
                side_a.push(i);
            } else {
                side_b.push(i);
            }
        }
    }

    // If one side is empty (no cut edges), fall back to midpoint split.
    if side_a.is_empty() || side_b.is_empty() {
        let mid = n / 2;
        side_a = (0..mid).collect();
        side_b = (mid..n).collect();
    }

    // The group with higher mean sensitivity becomes the "sensitive" group.
    let mean_of = |indices: &[usize]| -> f32 {
        if indices.is_empty() {
            return 0.0;
        }
        indices.iter().map(|&i| sensitivity[i]).sum::<f32>() / indices.len() as f32
    };

    if mean_of(&side_a) >= mean_of(&side_b) {
        (side_a, side_b)
    } else {
        (side_b, side_a)
    }
}

/// Convert a mincut partition into per-subcarrier importance weights.
///
/// Sensitive subcarriers (high body-motion correlation) get weight > 1.0,
/// insensitive ones get weight 0.5. This allows downstream feature extraction
/// to emphasise the most informative subcarriers.
pub fn subcarrier_importance_weights(sensitivity: &[f32]) -> Vec<f32> {
    if sensitivity.is_empty() {
        return vec![];
    }
    let (sensitive, _insensitive) = mincut_subcarrier_partition(sensitivity);
    let max_sens = sensitivity
        .iter()
        .cloned()
        .fold(f32::NEG_INFINITY, f32::max)
        .max(1e-9);

    let mut weights = vec![0.5f32; sensitivity.len()];
    for &idx in &sensitive {
        weights[idx] = 1.0 + (sensitivity[idx] / max_sens).min(1.0);
    }
    weights
}

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

    #[test]
    fn partition_covers_all_indices() {
        let sensitivity: Vec<f32> = (0..10).map(|i| i as f32 * 0.1).collect();
        let (sensitive, insensitive) = mincut_subcarrier_partition(&sensitivity);

        // Both groups must be non-empty for a non-trivial input.
        assert!(!sensitive.is_empty(), "sensitive group must not be empty");
        assert!(
            !insensitive.is_empty(),
            "insensitive group must not be empty"
        );

        // Together they must cover every index exactly once.
        let mut all_indices: Vec<usize> = sensitive
            .iter()
            .chain(insensitive.iter())
            .cloned()
            .collect();
        all_indices.sort_unstable();
        let expected: Vec<usize> = (0..10).collect();
        assert_eq!(all_indices, expected, "partition must cover all 10 indices");
    }

    #[test]
    fn partition_empty_input() {
        let (s, i) = mincut_subcarrier_partition(&[]);
        assert!(s.is_empty());
        assert!(i.is_empty());
    }

    #[test]
    fn partition_single_element() {
        let (s, i) = mincut_subcarrier_partition(&[0.5]);
        assert_eq!(s, vec![0]);
        assert!(i.is_empty());
    }

    #[test]
    fn test_importance_weights_empty() {
        let w = subcarrier_importance_weights(&[]);
        assert!(w.is_empty());
    }

    #[test]
    fn test_importance_weights_all_equal() {
        let sensitivity = vec![1.0f32; 8];
        let w = subcarrier_importance_weights(&sensitivity);
        assert_eq!(w.len(), 8);
        // All subcarriers have identical sensitivity so all should be classified
        // the same way (either all sensitive or all insensitive after mincut).
        // At minimum, no weight should exceed 2.0 or be negative.
        for &wt in &w {
            assert!((0.5..=2.0).contains(&wt), "weight {wt} out of range");
        }
    }

    #[test]
    fn test_importance_weights_sensitive_higher() {
        // First 5 subcarriers have high sensitivity, last 5 low.
        let sensitivity: Vec<f32> = (0..10).map(|i| if i < 5 { 0.9 } else { 0.1 }).collect();
        let w = subcarrier_importance_weights(&sensitivity);
        assert_eq!(w.len(), 10);

        let mean_high: f32 = w[..5].iter().sum::<f32>() / 5.0;
        let mean_low: f32 = w[5..].iter().sum::<f32>() / 5.0;
        assert!(
            mean_high > mean_low,
            "sensitive subcarriers should have higher mean weight ({mean_high}) than insensitive ({mean_low})"
        );
    }
}
