/// Cut tree at rk threshold to produce cluster assignments.
///
/// This implements a two-pass algorithm:
/// 1. Top-down: Force children of high-rk nodes to be above threshold
/// 2. Bottom-up: Collect cluster members
///
/// # Arguments
/// * `merges` - Merge matrix using R-style indexing (negative=leaf, positive=internal)
/// * `rk` - rk values for internal nodes only (length = n_merges)
/// * `threshold` - Cut threshold (e.g., ln(0.5))
/// * `n_isolates` - Number of leaf nodes
///
/// # Returns
/// Vector of cluster labels (1-indexed) for each isolate (length = n_isolates)
pub fn summarise_clusters(
    merges: &[(i32, i32)],
    rk: &[f64],
    threshold: f64,
    n_isolates: usize,
) -> Vec<usize> {
    let n_merges = merges.len();
    
    // Working copy of rk that we can modify
    let mut temp_rk = rk.to_vec();
    
    // Track which internal nodes get consumed by parents
    let mut keep = vec![false; n_merges];
    
    // Cluster member lists for each internal node
    let mut clusters: Vec<Vec<usize>> = vec![Vec::new(); n_merges];

    // Pass 1: Top-down propagation
    // If a node is above threshold, force its internal children above threshold too
    for i in (0..n_merges).rev() {
        if temp_rk[i] > threshold {
            // Left child
            if merges[i].0 > 0 {
                let child_idx = (merges[i].0 - 1) as usize;
                temp_rk[child_idx] = threshold + 1.0;
            }
            // Right child
            if merges[i].1 > 0 {
                let child_idx = (merges[i].1 - 1) as usize;
                temp_rk[child_idx] = threshold + 1.0;
            }
        }
    }

    // Pass 2: Bottom-up collection of cluster members
    for i in 0..n_merges {
        if temp_rk[i] > threshold {
            // Process left child
            if merges[i].0 < 0 {
                let leaf_id = merges[i].0.abs() as usize;
                clusters[i].push(leaf_id);
            } else {
                let child_idx = (merges[i].0 - 1) as usize;
                let drained: Vec<usize> = clusters[child_idx].drain(..).collect();
                clusters[i].extend(drained);
                keep[child_idx] = true;
            }

            // Process right child
            if merges[i].1 < 0 {
                let leaf_id = merges[i].1.abs() as usize;
                clusters[i].push(leaf_id);
            } else {
                let child_idx = (merges[i].1 - 1) as usize;
                let drained: Vec<usize> = clusters[child_idx].drain(..).collect();
                clusters[i].extend(drained);
                keep[child_idx] = true;
            }
        }
    }

    // Final: Assign cluster labels
    let mut final_clust = vec![0usize; n_isolates];
    let mut n_clust = 0;

    // Assign labels to clusters that weren't consumed
    for i in 0..n_merges {
        if !keep[i] && !clusters[i].is_empty() {
            n_clust += 1;
            for &member in &clusters[i] {
                final_clust[member - 1] = n_clust; // Convert to 0-indexed
            }
        }
    }

    // Assign singleton clusters to any unclustered isolates
    for i in 0..n_isolates {
        if final_clust[i] == 0 {
            n_clust += 1;
            final_clust[i] = n_clust;
        }
    }

    final_clust
}

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

    #[test]
    fn test_summarise_clusters_all_separate() {
        // 3-leaf tree where all nodes are below threshold
        let merges = vec![
            (-1, -2),
            (1, -3),
        ];
        let rk = vec![-1.0, -1.0]; // All below threshold
        let threshold = 0.0;
        let n_isolates = 3;

        let result = summarise_clusters(&merges, &rk, threshold, n_isolates);

        // All isolates should be in separate clusters
        assert_eq!(result.len(), 3);
        assert_ne!(result[0], result[1]);
        assert_ne!(result[1], result[2]);
        assert_ne!(result[0], result[2]);
    }

    #[test]
    fn test_summarise_clusters_all_together() {
        // 3-leaf tree where root is above threshold
        let merges = vec![
            (-1, -2),
            (1, -3),
        ];
        let rk = vec![-1.0, 1.0]; // Root above threshold
        let threshold = 0.0;
        let n_isolates = 3;

        let result = summarise_clusters(&merges, &rk, threshold, n_isolates);

        // All isolates should be in the same cluster
        assert_eq!(result.len(), 3);
        assert_eq!(result[0], result[1]);
        assert_eq!(result[1], result[2]);
    }

    #[test]
    fn test_summarise_clusters_partial() {
        // 4-leaf tree: ((1,2),(3,4))
        let merges = vec![
            (-1, -2), // merge 1,2
            (-3, -4), // merge 3,4
            (1, 2),   // merge the two pairs
        ];
        let rk = vec![1.0, 1.0, -1.0]; // First two merges above, root below
        let threshold = 0.0;
        let n_isolates = 4;

        let result = summarise_clusters(&merges, &rk, threshold, n_isolates);

        // Should have 2 clusters: {1,2} and {3,4}
        assert_eq!(result.len(), 4);
        assert_eq!(result[0], result[1]); // 1 and 2 together
        assert_eq!(result[2], result[3]); // 3 and 4 together
        assert_ne!(result[0], result[2]); // But different clusters
    }
}
