//! Pre-computed lgamma lookup table for fast Bayesian hierarchical clustering
//!
//! Ported from fastbaps C++ implementation (tree_llk.cpp, bhier_parallel.cpp)
//! The table stores lgamma(count + prior_idx * 0.001) values for efficient lookup
//! during likelihood computations.

use ndarray::Array2;
use statrs::function::gamma::ln_gamma;

/// Pre-computed lgamma lookup table
///
/// Stores lgamma values quantized to 0.001 precision to avoid repeated
/// expensive lgamma() calls during Bayesian hierarchical clustering.
///
/// The table is indexed by:
/// - row: count (0..n_isolates+2)
/// - col: prior_idx representing prior/0.001 rounded up
pub struct LgammaTable {
    /// 2D table stored in row-major order
    data: Vec<f64>,
    n_rows: usize,
    n_cols: usize,
}

impl LgammaTable {
    /// Build a new lgamma lookup table
    ///
    /// # Arguments
    /// * `n_isolates` - Number of isolates in the dataset
    /// * `max_prior` - Maximum prior value across all SNPs and bases
    ///
    /// # Returns
    /// Pre-computed table covering all needed lgamma values
    pub fn new(n_isolates: usize, max_prior: f64) -> Self {
        let lg_size = (max_prior * 1000.0).ceil() as usize + 3;
        let n_rows = n_isolates + 2;
        let n_cols = lg_size;
        
        let mut data = vec![0.0; n_rows * n_cols];
        
        for i in 0..n_rows {
            for j in 0..n_cols {
                // Skip (0,0) as in original C++ code
                if i == 0 && j == 0 {
                    continue;
                }
                
                let value = i as f64 + j as f64 * 0.001;
                data[i * n_cols + j] = ln_gamma(value);
            }
        }
        
        Self {
            data,
            n_rows,
            n_cols,
        }
    }
    
    /// Lookup lgamma value at (count, prior_idx)
    ///
    /// # Arguments
    /// * `count` - Count value (row index)
    /// * `prior_idx` - Prior index (column index)
    ///
    /// # Returns
    /// Pre-computed lgamma(count + prior_idx * 0.001)
    #[inline]
    pub fn get(&self, count: usize, prior_idx: usize) -> f64 {
        debug_assert!(count < self.n_rows, "count out of bounds");
        debug_assert!(prior_idx < self.n_cols, "prior_idx out of bounds");
        self.data[count * self.n_cols + prior_idx]
    }
    
    /// Get number of rows (n_isolates + 2)
    #[inline]
    pub fn n_rows(&self) -> usize {
        self.n_rows
    }
    
    /// Get number of columns (lg_size)
    #[inline]
    pub fn n_cols(&self) -> usize {
        self.n_cols
    }
}

/// Convert prior matrix to index matrix for lgamma table lookup
///
/// Each prior value is converted to an index via ceil(prior / 0.001)
///
/// # Arguments
/// * `prior` - Prior matrix (5 x n_snps), rows = A/C/G/T/gap
///
/// # Returns
/// Index matrix (5 x n_snps) for lgamma table lookup
pub fn compute_prior_index(prior: &Array2<f64>) -> Array2<usize> {
    prior.mapv(|val| (val / 0.001).ceil() as usize)
}

/// Compute term1 vector for likelihood calculation
///
/// Port of C++ code:
/// ```cpp
/// for (partition_length=1; partition_length<n_isolates+2; partition_length++){
///     for (j=0; j < n_snps; j++){
///         alpha_sum = 0.0;
///         for(i=0; i<5; i++){
///             alpha_sum += prior(i,j);
///         }
///         term1(partition_length) += lgamma(alpha_sum) - lgamma(partition_length+alpha_sum);
///     }
/// }
/// ```
///
/// # Arguments
/// * `n_isolates` - Number of isolates
/// * `n_snps` - Number of SNPs
/// * `prior` - Prior matrix (5 x n_snps)
///
/// # Returns
/// Vector of length n_isolates+2 with term1 values
pub fn compute_term1(n_isolates: usize, n_snps: usize, prior: &Array2<f64>) -> Vec<f64> {
    let mut term1 = vec![0.0; n_isolates + 2];
    
    for partition_length in 1..(n_isolates + 2) {
        for j in 0..n_snps {
            // Sum prior values across all 5 bases for this SNP
            let mut alpha_sum = 0.0;
            for i in 0..5 {
                alpha_sum += prior[[i, j]];
            }
            
            // Compute lgamma(alpha_sum) - lgamma(partition_length + alpha_sum)
            let lg_alpha = ln_gamma(alpha_sum);
            let lg_sum = ln_gamma(partition_length as f64 + alpha_sum);
            term1[partition_length] += lg_alpha - lg_sum;
        }
    }
    
    term1
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_lgamma_table_creation() {
        let table = LgammaTable::new(10, 5.0);
        assert_eq!(table.n_rows(), 12); // 10 + 2
        assert_eq!(table.n_cols(), 5003); // ceil(5.0 * 1000) + 3
    }
    
    #[test]
    fn test_lgamma_table_lookup() {
        let table = LgammaTable::new(10, 5.0);
        
        // Test a few values manually
        let val = table.get(1, 1);
        let expected = ln_gamma(1.0 + 0.001);
        assert!((val - expected).abs() < 1e-10);
        
        let val = table.get(5, 100);
        let expected = ln_gamma(5.0 + 0.1);
        assert!((val - expected).abs() < 1e-10);
    }
    
    #[test]
    fn test_compute_prior_index() {
        let prior = Array2::from_shape_vec((5, 3), vec![
            0.001, 0.002, 0.003,
            0.010, 0.020, 0.030,
            0.100, 0.200, 0.300,
            1.000, 2.000, 3.000,
            0.500, 1.500, 2.500,
        ]).unwrap();
        
        let index = compute_prior_index(&prior);
        
        assert_eq!(index[[0, 0]], 1);  // ceil(0.001 / 0.001) = 1
        assert_eq!(index[[0, 1]], 2);  // ceil(0.002 / 0.001) = 2
        assert_eq!(index[[1, 0]], 10); // ceil(0.010 / 0.001) = 10
        assert_eq!(index[[3, 0]], 1000); // ceil(1.000 / 0.001) = 1000
    }
    
    #[test]
    fn test_compute_term1() {
        let prior = Array2::from_shape_vec((5, 2), vec![
            1.0, 2.0,
            1.0, 2.0,
            1.0, 2.0,
            1.0, 2.0,
            1.0, 2.0,
        ]).unwrap();
        
        let term1 = compute_term1(3, 2, &prior);
        
        assert_eq!(term1.len(), 5); // 3 + 2
        assert_eq!(term1[0], 0.0);  // Not computed for partition_length=0
        
        // For partition_length=1, each SNP contributes:
        // lgamma(5.0) - lgamma(1.0 + 5.0) for SNP 0
        // lgamma(10.0) - lgamma(1.0 + 10.0) for SNP 1
        let expected = (ln_gamma(5.0) - ln_gamma(6.0)) +
                      (ln_gamma(10.0) - ln_gamma(11.0));
        assert!((term1[1] - expected).abs() < 1e-10);
    }
}
