//! Tree node likelihood computation for Bayesian hierarchical clustering.
//!
//! This module traverses a merge tree bottom-up, computing dk, mllk, ptree, and rk
//! for each node. The algorithm follows the C++ implementation in tree_llk.cpp.

use crate::lgamma_table::{compute_prior_index, compute_term1, LgammaTable};
use crate::likelihood::marginal_llk;
use crate::logspace::log_sum_exp;
use ndarray::Array2;
use sprs::{CsMat, TriMat};

/// Result of tree likelihood computation
#[derive(Debug)]
pub struct TreeLlkResult {
    /// dk[i] = log Dirichlet Process prior mass for node i
    pub dk: Vec<f64>,
    /// mllk[i] = marginal log-likelihood of node i as a single cluster
    pub mllk: Vec<f64>,
    /// rk[i] = log-ratio statistic for splitting node i
    pub rk: Vec<f64>,
    /// ptree[i] = log probability of subtree rooted at node i
    pub ptree: Vec<f64>,
    /// clust_size[i] = number of isolates in node i
    pub clust_size: Vec<usize>,
}

/// Compute tree node likelihoods for all nodes in a hierarchical clustering.
///
/// # Algorithm
/// 1. Pre-compute lgamma table, prior_index, term1
/// 2. Count initial alleles for each isolate into sparse count matrices
/// 3. Initialize leaf nodes (i=0..n_isolates):
///    - clust_size[i] = 1, dk[i] = alpha = 0.0
///    - mllk[i] = marginal_llk for single-isolate partition
///    - ptree[i] = mllk[i]
///    - rk[i] = +infinity (leaves cannot be split)
/// 4. For each merge m (i = m + n_isolates):
///    - Decode left/right children from R-style merge matrix
///    - clust_size[i] = clust_size[left] + clust_size[right]
///    - sp_partition_counts[i] = sp_partition_counts[left] + sp_partition_counts[right]
///    - dk[i] = log_sum_exp(lgamma(n), dk[left] + dk[right])
///    - mllk[i] = marginal_llk(sp_partition_counts[i], clust_size[i], ...)
///    - ptree[i] = log_sum_exp(lg_table.get(n,0) + mllk[i] - dk[i],
///                             dk[left] + dk[right] + ptree[left] + ptree[right] - dk[i])
///    - rk[i] = mllk[i] - log_sum_exp(mllk[i], dk[left] + dk[right] - lg_table.get(n,0) + ptree[left] + ptree[right])
///
/// # Arguments
/// * `snp_data` - CSC data array (values 1-4 for A/C/G/T)
/// * `snp_indices` - CSC row indices (SNP positions)
/// * `snp_indptr` - CSC column pointers (isolate boundaries)
/// * `n_snps` - Number of SNP positions
/// * `n_isolates` - Number of isolates
/// * `consensus` - Consensus sequence (not used in computation, kept for interface consistency)
/// * `prior` - 5 x n_snps prior matrix (rows: consensus/A/C/G/T)
/// * `merges` - R-style merge matrix: negative = leaf (-i means isolate i-1), positive = internal node (j means node n_isolates+j-1)
///
/// # Returns
/// TreeLlkResult with dk, mllk, rk, ptree, and clust_size vectors
pub fn tree_llk(
    snp_data: &[u8],
    snp_indices: &[usize],
    snp_indptr: &[usize],
    n_snps: usize,
    n_isolates: usize,
    _consensus: &[f64],
    prior: &Array2<f64>,
    merges: &[(i32, i32)],
) -> TreeLlkResult {
    // Alpha = log(1) = 0.0 in original C++ code
    let alpha = 0.0;

    let n_merges = merges.len();
    let n_nodes = n_merges + n_isolates;

    // Pre-compute lgamma table
    let max_prior = prior.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
    let lg_table = LgammaTable::new(n_isolates, max_prior);

    // Pre-compute prior_index matrix
    let prior_index = compute_prior_index(prior);

    // Pre-compute term1 vector
    let term1 = compute_term1(n_isolates, n_snps, prior);

    // Count initial alleles for each isolate into sparse count matrices
    let mut sp_partition_counts: Vec<CsMat<u32>> = Vec::with_capacity(n_nodes);

    for isolate_idx in 0..n_isolates {
        let start = snp_indptr[isolate_idx];
        let end = snp_indptr[isolate_idx + 1];

        let mut trimat: TriMat<u32> = TriMat::new((n_snps, 4));

        for i in start..end {
            let snp_idx = snp_indices[i];
            let allele_value = snp_data[i]; // 1-4 for A/C/G/T
            let allele_idx = (allele_value - 1) as usize; // 0-3 for matrix column

            trimat.add_triplet(snp_idx, allele_idx, 1);
        }

        sp_partition_counts.push(trimat.to_csr());
    }

    // Initialize result vectors
    let mut clust_size = vec![0; n_nodes];
    let mut dk = vec![0.0; n_nodes];
    let mut mllk = vec![0.0; n_nodes];
    let mut ptree = vec![0.0; n_nodes];
    let mut rk = vec![0.0; n_nodes];

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

        // Compute marginal likelihood for single isolate
        mllk[i] = marginal_llk(
            &sp_partition_counts[i],
            1,
            n_snps,
            &term1,
            &lg_table,
            &prior_index,
        );

        ptree[i] = mllk[i];
        rk[i] = f64::INFINITY;
    }

    // Process each merge
    for m in 0..n_merges {
        let i = m + n_isolates;

        // Decode left and right children from R-style merge matrix
        let left = if merges[m].0 < 0 {
            // Negative = leaf: -1 means isolate 0, -2 means isolate 1, etc.
            (merges[m].0.abs() - 1) as usize
        } else {
            // Positive = internal node: 1 means first merge (index n_isolates), etc.
            n_isolates + (merges[m].0 - 1) as usize
        };

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

        // Merge cluster sizes
        clust_size[i] = clust_size[left] + clust_size[right];
        let n = clust_size[i];

        // Merge sparse count matrices
        let sp_merged = &sp_partition_counts[left] + &sp_partition_counts[right];
        sp_partition_counts.push(sp_merged);

        // Compute dk[i] = log_sum_exp(alpha + lgamma(n), dk[left] + dk[right])
        // Since alpha = 0.0, this becomes: log_sum_exp(lgamma(n), dk[left] + dk[right])
        // lgamma(n) = lg_table.get(n, 0)
        dk[i] = log_sum_exp(lg_table.get(n, 0), dk[left] + dk[right]);

        // Compute marginal likelihood for merged node
        mllk[i] = marginal_llk(
            &sp_partition_counts[i],
            n,
            n_snps,
            &term1,
            &lg_table,
            &prior_index,
        );

        // Compute ptree[i]
        // ptree[i] = log_sum_exp(alpha + lgamma(n) + mllk[i] - dk[i],
        //                        dk[left] + dk[right] + ptree[left] + ptree[right] - dk[i])
        // With alpha = 0.0:
        let term_a = lg_table.get(n, 0) + mllk[i] - dk[i];
        let term_b = dk[left] + dk[right] + ptree[left] + ptree[right] - dk[i];
        ptree[i] = log_sum_exp(term_a, term_b);

        // Compute rk[i]
        // rk[i] = mllk[i] - log_sum_exp(mllk[i], dk[left] + dk[right] - alpha - lgamma(n) + ptree[left] + ptree[right])
        // With alpha = 0.0:
        let denom = log_sum_exp(
            mllk[i],
            dk[left] + dk[right] - lg_table.get(n, 0) + ptree[left] + ptree[right],
        );
        rk[i] = mllk[i] - denom;
    }

    TreeLlkResult {
        dk,
        mllk,
        rk,
        ptree,
        clust_size,
    }
}

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

    #[test]
    fn test_tree_llk_single_isolate() {
        // Test with a single isolate (no merges)
        let n_snps = 3;
        let n_isolates = 1;

        // Empty SNP data (all consensus)
        let snp_data: Vec<u8> = vec![];
        let snp_indices: Vec<usize> = vec![];
        let snp_indptr: Vec<usize> = vec![0, 0]; // One isolate, no non-zero entries

        let consensus = vec![1.0; n_snps];
        let prior = Array2::ones((5, n_snps));

        let merges: Vec<(i32, i32)> = vec![];

        let result = tree_llk(
            &snp_data,
            &snp_indices,
            &snp_indptr,
            n_snps,
            n_isolates,
            &consensus,
            &prior,
            &merges,
        );

        assert_eq!(result.clust_size.len(), 1);
        assert_eq!(result.clust_size[0], 1);
        assert_eq!(result.dk[0], 0.0); // alpha = 0
        assert!(result.mllk[0].is_finite());
        assert_eq!(result.ptree[0], result.mllk[0]);
        assert_eq!(result.rk[0], f64::INFINITY);
    }

    #[test]
    fn test_tree_llk_two_isolates_one_merge() {
        // Test with two isolates and one merge
        let n_snps = 2;
        let n_isolates = 2;

        // Isolate 0: SNP 0 has allele 1 (A)
        // Isolate 1: SNP 1 has allele 2 (C)
        let snp_data: Vec<u8> = vec![1, 2];
        let snp_indices: Vec<usize> = vec![0, 1];
        let snp_indptr: Vec<usize> = vec![0, 1, 2];

        let consensus = vec![1.0; n_snps];
        let prior = Array2::ones((5, n_snps));

        // One merge: combine isolates 0 and 1
        // R-style: negative indices for leaves (-1 = isolate 0, -2 = isolate 1)
        let merges: Vec<(i32, i32)> = vec![(-1, -2)];

        let result = tree_llk(
            &snp_data,
            &snp_indices,
            &snp_indptr,
            n_snps,
            n_isolates,
            &consensus,
            &prior,
            &merges,
        );

        // Should have 3 nodes: 2 leaves + 1 internal
        assert_eq!(result.clust_size.len(), 3);
        assert_eq!(result.clust_size[0], 1); // leaf 0
        assert_eq!(result.clust_size[1], 1); // leaf 1
        assert_eq!(result.clust_size[2], 2); // merged node

        // Leaves have rk = infinity
        assert_eq!(result.rk[0], f64::INFINITY);
        assert_eq!(result.rk[1], f64::INFINITY);

        // Merged node should have finite rk
        assert!(result.rk[2].is_finite());
        assert!(result.mllk[2].is_finite());
        assert!(result.ptree[2].is_finite());
    }

    #[test]
    fn test_tree_llk_three_isolates_two_merges() {
        // Test with three isolates and two merges
        let n_snps = 3;
        let n_isolates = 3;

        // All isolates have consensus (empty sparse data)
        let snp_data: Vec<u8> = vec![];
        let snp_indices: Vec<usize> = vec![];
        let snp_indptr: Vec<usize> = vec![0, 0, 0, 0];

        let consensus = vec![1.0; n_snps];
        let prior = Array2::ones((5, n_snps));

        // Two merges:
        // Merge 1: combine isolates 0 and 1 -> node 3
        // Merge 2: combine node 3 and isolate 2 -> node 4
        let merges: Vec<(i32, i32)> = vec![(-1, -2), (1, -3)];

        let result = tree_llk(
            &snp_data,
            &snp_indices,
            &snp_indptr,
            n_snps,
            n_isolates,
            &consensus,
            &prior,
            &merges,
        );

        // Should have 5 nodes: 3 leaves + 2 internal
        assert_eq!(result.clust_size.len(), 5);
        assert_eq!(result.clust_size[0], 1);
        assert_eq!(result.clust_size[1], 1);
        assert_eq!(result.clust_size[2], 1);
        assert_eq!(result.clust_size[3], 2); // first merge
        assert_eq!(result.clust_size[4], 3); // second merge

        // All internal nodes should have finite values
        for i in 3..5 {
            assert!(result.dk[i].is_finite());
            assert!(result.mllk[i].is_finite());
            assert!(result.ptree[i].is_finite());
            assert!(result.rk[i].is_finite());
        }
    }
}
