//! Marginal log-likelihood kernel for partition evaluation.
//!
//! This module implements the core likelihood computation used by both
//! tree_llk and Bayesian hierarchical clustering. It computes the marginal
//! log-likelihood of a partition given sparse allele counts.

use crate::lgamma_table::LgammaTable;
use ndarray::Array2;
use sprs::CsMat;

/// Compute marginal log-likelihood for a single partition.
///
/// # Arguments
/// * `sp_counts` - Sparse matrix (n_snps x 4) of allele counts in CSR format
/// * `partition_size` - Number of isolates in this partition
/// * `n_snps` - Number of SNP positions
/// * `term1` - Precomputed term1 vector indexed by partition size
/// * `lg_table` - Lgamma lookup table for fast lgamma evaluation
/// * `prior_index` - 5 x n_snps matrix of prior indices (row 0 = consensus, rows 1-4 = A/C/G/T)
///
/// # Formula
/// For each partition:
/// - Start with term1[partition_size]
/// - For each non-zero (snp, allele) with count c:
///   add lgamma(c + prior[1+allele, snp]) - lgamma(prior[1+allele, snp])
/// - For each SNP compute consensus_count = partition_size - sum_of_non_consensus_counts
/// - Add lgamma(consensus_count + prior[0, snp]) - lgamma(prior[0, snp])
pub fn marginal_llk(
    sp_counts: &CsMat<u32>,
    partition_size: usize,
    n_snps: usize,
    term1: &[f64],
    lg_table: &LgammaTable,
    prior_index: &Array2<usize>,
) -> f64 {
    debug_assert_eq!(sp_counts.rows(), n_snps);
    debug_assert_eq!(sp_counts.cols(), 4);
    debug_assert_eq!(prior_index.nrows(), 5);
    debug_assert_eq!(prior_index.ncols(), n_snps);

    // Initialize consensus counts to partition size (will subtract non-consensus)
    let mut consensus_counts = vec![partition_size; n_snps];
    
    // Start with term1 contribution
    let mut mllk = term1[partition_size];

    // Process non-zero allele counts
    // In sparse matrix: rows = SNP indices, cols = allele indices (0-3 for A/C/G/T)
    for (count, (snp_idx, allele_idx)) in sp_counts.iter() {
        let count = *count as usize;
        
        // Subtract from consensus
        consensus_counts[snp_idx] -= count;
        
        // Add lgamma(count + prior_idx) - lgamma(prior_idx) for this allele
        // prior_index row 0 = consensus, rows 1-4 = A/C/G/T alleles
        let prior_idx = prior_index[[1 + allele_idx, snp_idx]];
        mllk += lg_table.get(count, prior_idx) - lg_table.get(0, prior_idx);
    }

    // Add consensus contribution for each SNP
    for snp_idx in 0..n_snps {
        let consensus_count = consensus_counts[snp_idx];
        let prior_idx = prior_index[[0, snp_idx]];
        mllk += lg_table.get(consensus_count, prior_idx) - lg_table.get(0, prior_idx);
    }

    mllk
}

/// Compute marginal log-likelihood for two merged partitions.
///
/// This is used during hierarchical clustering to evaluate merge proposals
/// without explicitly constructing the merged sparse matrix.
///
/// # Arguments
/// * `sp_counts_a` - Sparse counts for partition A
/// * `sp_counts_b` - Sparse counts for partition B
/// * `size_a` - Number of isolates in partition A
/// * `size_b` - Number of isolates in partition B
/// * `n_snps` - Number of SNP positions
/// * `term1` - Precomputed term1 vector
/// * `lg_table` - Lgamma lookup table
/// * `prior_index` - 5 x n_snps matrix of prior indices
pub fn marginal_llk_merged(
    sp_counts_a: &CsMat<u32>,
    sp_counts_b: &CsMat<u32>,
    size_a: usize,
    size_b: usize,
    n_snps: usize,
    term1: &[f64],
    lg_table: &LgammaTable,
    prior_index: &Array2<usize>,
) -> f64 {
    debug_assert_eq!(sp_counts_a.rows(), n_snps);
    debug_assert_eq!(sp_counts_b.rows(), n_snps);
    debug_assert_eq!(sp_counts_a.cols(), 4);
    debug_assert_eq!(sp_counts_b.cols(), 4);

    // Merge the two sparse matrices
    let merged_size = size_a + size_b;
    let sp_counts_merged = sp_counts_a + sp_counts_b;

    // Use the single-partition kernel
    marginal_llk(
        &sp_counts_merged,
        merged_size,
        n_snps,
        term1,
        lg_table,
        prior_index,
    )
}

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

    #[test]
    fn test_marginal_llk_empty_partition() {
        // Empty partition (all consensus)
        let n_snps = 3;
        let partition_size = 10;
        
        // Create empty sparse matrix (no non-consensus alleles) using TriMat
        let trimat: TriMat<u32> = TriMat::new((n_snps, 4));
        let sp_counts = trimat.to_csr();
        
        // Mock term1 and lgamma table
        let term1 = vec![0.0; 100];
        let lg_table = LgammaTable::new(100, 1.0);
        
        // Mock prior_index (all ones to avoid index 0 issues)
        let prior_index = Array2::ones((5, n_snps));
        
        let mllk = marginal_llk(&sp_counts, partition_size, n_snps, &term1, &lg_table, &prior_index);
        
        // For all consensus: term1[10] + 3 * (lgamma(10+prior) - lgamma(prior))
        // This is just a sanity check that it runs
        assert!(mllk.is_finite());
    }

    #[test]
    fn test_marginal_llk_merged() {
        let n_snps = 2;
        
        // Create two simple sparse matrices with proper shape using TriMat
        let trimat_a: TriMat<u32> = TriMat::new((n_snps, 4));
        let sp_a = trimat_a.to_csr();
        
        let trimat_b: TriMat<u32> = TriMat::new((n_snps, 4));
        let sp_b = trimat_b.to_csr();
        
        let term1 = vec![0.0; 100];
        let lg_table = LgammaTable::new(100, 1.0);
        let prior_index = Array2::ones((5, n_snps));
        
        let mllk = marginal_llk_merged(&sp_a, &sp_b, 5, 5, n_snps, &term1, &lg_table, &prior_index);
        
        assert!(mllk.is_finite());
    }
}
