use crate::logspace::log_sum_exp;
use statrs::function::gamma::ln_gamma;

/// Compute dk and cluster sizes from a merge matrix.
///
/// This implements the bottom-up calculation of dk (marginal likelihood terms)
/// and cluster sizes for a hierarchical clustering tree.
///
/// # Arguments
/// * `merges` - Merge matrix (n_merges x 2) using R-style indexing:
///              - Negative values: leaf nodes (1-indexed, abs value is leaf ID)
///              - Positive values: internal nodes (1-indexed, offset by n_isolates)
/// * `n_isolates` - Number of leaf nodes
///
/// # Returns
/// Tuple of (dk, clust_sizes) where:
/// - dk: Vector of length n_nodes with log marginal likelihood terms
/// - clust_sizes: Vector of length n_nodes with cluster sizes
///
/// Both vectors have indices [0..n_isolates) for leaves and [n_isolates..n_nodes) for internal nodes.
pub fn calc_ddk(merges: &[(i32, i32)], n_isolates: usize) -> (Vec<f64>, Vec<usize>) {
    let n_merges = merges.len();
    let n_nodes = n_merges + n_isolates;

    let mut clust_size = vec![0usize; n_nodes];
    let mut dk = vec![0.0f64; n_nodes];

    // Initialize leaf nodes
    for i in 0..n_isolates {
        clust_size[i] = 1;
        dk[i] = 0.0; // log(1) = 0
    }

    // Process merges bottom-up
    for m in 0..n_merges {
        let i = m + n_isolates;

        // Decode left child (R-style indexing)
        let left = if merges[m].0 < 0 {
            (merges[m].0.abs() - 1) as usize
        } else {
            n_isolates + (merges[m].0 - 1) as usize
        };

        // Decode right child (R-style indexing)
        let right = if merges[m].1 < 0 {
            (merges[m].1.abs() - 1) as usize
        } else {
            n_isolates + (merges[m].1 - 1) as usize
        };

        clust_size[i] = clust_size[left] + clust_size[right];
        
        // dk[i] = log_sum_exp(lgamma(clust_size[i]), dk[left] + dk[right])
        let lgamma_n = ln_gamma(clust_size[i] as f64);
        dk[i] = log_sum_exp(lgamma_n, dk[left] + dk[right]);
    }

    (dk, clust_size)
}

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

    #[test]
    fn test_calc_ddk_simple() {
        // Simple 3-leaf tree: merge (1,2) then merge that with 3
        // R-style: leaves are negative 1-indexed
        let merges = vec![
            (-1, -2), // merge leaves 1 and 2
            (1, -3),  // merge internal node 1 with leaf 3
        ];
        let n_isolates = 3;

        let (dk, clust_size) = calc_ddk(&merges, n_isolates);

        // Check sizes
        assert_eq!(dk.len(), 5); // 3 leaves + 2 internal
        assert_eq!(clust_size.len(), 5);

        // Leaf sizes should be 1
        assert_eq!(clust_size[0], 1);
        assert_eq!(clust_size[1], 1);
        assert_eq!(clust_size[2], 1);

        // First merge: size 2
        assert_eq!(clust_size[3], 2);

        // Second merge: size 3
        assert_eq!(clust_size[4], 3);

        // Leaf dk should be 0
        assert_eq!(dk[0], 0.0);
        assert_eq!(dk[1], 0.0);
        assert_eq!(dk[2], 0.0);
    }
}
