//! Bayesian hierarchical clustering engine with Rayon parallelism.
//!
//! Ported from fastbaps C++ implementation (bhier_parallel.cpp).
//! Implements greedy agglomerative clustering guided by marginal likelihood.

use crate::lgamma_table::LgammaTable;
use crate::likelihood::marginal_llk;
use crate::logspace::log_sum_exp;
use ndarray::Array2;
use rayon::prelude::*;
use sprs::{CsMat, TriMat};

/// Result of Bayesian hierarchical clustering
#[derive(Debug, Clone)]
pub struct BhcResult {
    /// Merge edges (R-style 1-indexed), one per merge step
    /// Negative values = original partitions, positive = internal nodes
    pub edges: Vec<(i32, i32)>,
    
    /// Marginal log-likelihood at each merge (heights for dendrogram)
    pub heights: Vec<f64>,
    
    /// rk values at each merge (posterior probability metric)
    pub rk: Vec<f64>,
    
    /// Initial marginal log-likelihood for each starting partition
    pub initial_partition_llk: Vec<f64>,
}

/// Bayesian hierarchical clustering with Rayon parallelism.
///
/// Performs greedy agglomerative clustering on partitions, merging pairs
/// with the highest rk (posterior probability of belonging together).
///
/// # Algorithm
/// 1. Build sparse count matrices for each initial partition
/// 2. Compute initial all-pairs rk values (parallel)
/// 3. Iteratively merge highest rk pair, update rk matrix (parallel update)
///
/// # Arguments
/// * `snp_data` - CSC sparse matrix data (values 1-4 for A/C/G/T)
/// * `snp_indices` - CSC row indices
/// * `snp_indptr` - CSC column pointers
/// * `n_snps` - Number of SNP positions (rows)
/// * `n_isolates` - Number of isolates (columns)
/// * `prior` - Prior matrix (5 x n_snps), rows = consensus/A/C/G/T
/// * `partitions` - Initial partitions (list of 1-indexed isolate IDs)
/// * `dk_initial` - Initial dk values for each partition
/// * `n_cores` - Number of threads for Rayon
///
/// # Returns
/// BhcResult containing merge tree and likelihoods
pub fn bhier_parallel(
    snp_data: &[u8],
    snp_indices: &[usize],
    snp_indptr: &[usize],
    n_snps: usize,
    n_isolates: usize,
    prior: &Array2<f64>,
    partitions: &[Vec<usize>],
    dk_initial: &[f64],
    n_cores: usize,
) -> BhcResult {
    // Confine all parallelism to a scoped pool sized for THIS call. The previous
    // build_global() only ever took effect on the first invocation in a process,
    // so a later call with a different n_cores -- or any other crate seeding the
    // global pool first -- would silently keep the wrong thread count. A scoped
    // pool plus install() guarantees every par_iter below runs on exactly
    // n_cores threads and is torn down when the call returns.
    let pool = rayon::ThreadPoolBuilder::new()
        .num_threads(n_cores)
        .build()
        .expect("failed to build Rayon thread pool");

    pool.install(|| {
        bhier_parallel_inner(
            snp_data,
            snp_indices,
            snp_indptr,
            n_snps,
            n_isolates,
            prior,
            partitions,
            dk_initial,
        )
    })
}

fn bhier_parallel_inner(
    snp_data: &[u8],
    snp_indices: &[usize],
    snp_indptr: &[usize],
    n_snps: usize,
    n_isolates: usize,
    prior: &Array2<f64>,
    partitions: &[Vec<usize>],
    dk_initial: &[f64],
) -> BhcResult {
    let n_partitions = partitions.len();
    
    // Build lgamma lookup table and precomputed terms
    let max_prior = prior.iter().copied().fold(f64::NEG_INFINITY, f64::max);
    let lg_table = LgammaTable::new(n_isolates, max_prior);
    let prior_index = crate::lgamma_table::compute_prior_index(prior);
    let term1 = crate::lgamma_table::compute_term1(n_isolates, n_snps, prior);

    // Phase 1: Build sparse count matrices for each partition
    let mut sp_partition_counts: Vec<CsMat<u32>> = partitions
        .par_iter()
        .map(|partition| build_partition_counts(
            partition,
            snp_data,
            snp_indices,
            snp_indptr,
            n_snps,
        ))
        .collect();

    let mut partition_sizes: Vec<usize> = partitions.iter().map(|p| p.len()).collect();

    // Compute initial partition likelihoods (parallel)
    let initial_partition_llk: Vec<f64> = sp_partition_counts
        .par_iter()
        .enumerate()
        .map(|(i, counts)| {
            marginal_llk(
                counts,
                partition_sizes[i],
                n_snps,
                &term1,
                &lg_table,
                &prior_index,
            )
        })
        .collect();

    // Initialize ptree with initial likelihoods
    let mut ptree = initial_partition_llk.clone();
    let mut dk = dk_initial.to_vec();

    // Phase 2: Compute initial all-pairs rk values (parallel over outer loop)
    let n_pairs = n_partitions * n_partitions;
    let rk_init: Vec<f64> = (0..n_pairs)
        .into_par_iter()
        .map(|idx| {
            let p1 = idx / n_partitions;
            let p2 = idx % n_partitions;
            
            if p1 >= p2 {
                return f64::NEG_INFINITY;
            }
            
            // Compute merged likelihood
            let merged = &sp_partition_counts[p1] + &sp_partition_counts[p2];
            let size_merged = partition_sizes[p1] + partition_sizes[p2];
            let mllk_merged = marginal_llk(
                &merged,
                size_merged,
                n_snps,
                &term1,
                &lg_table,
                &prior_index,
            );
            
            // rk = mllk_merged - ptree[p1] - ptree[p2] + lgamma(size_merged) - dk[p1] - dk[p2]
            let rk_val = mllk_merged - ptree[p1] - ptree[p2] 
                + lg_table.get(size_merged, 0) - dk[p1] - dk[p2];
            
            rk_val
        })
        .collect();

    let mut rk_matrix = rk_init;
    // Fill symmetric part
    for p1 in 0..n_partitions {
        for p2 in (p1 + 1)..n_partitions {
            let val = rk_matrix[p1 * n_partitions + p2];
            rk_matrix[p2 * n_partitions + p1] = val;
        }
    }

    // Phase 3: Greedy agglomeration loop
    let mut edges = Vec::with_capacity(n_partitions - 1);
    let mut heights = Vec::with_capacity(n_partitions - 1);
    let mut rk_vec = Vec::with_capacity(n_partitions - 1);
    let mut group_mems: Vec<i32> = (1..=n_partitions as i32).map(|x| -x).collect();
    let mut used = vec![false; n_partitions];

    for merge_step in 0..(n_partitions - 1) {
        // Find max rk
        let (max_idx, &max_rk) = rk_matrix
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .unwrap();
        
        let row_index = max_idx / n_partitions;
        let col_index = max_idx % n_partitions;
        
        let min_i = row_index.min(col_index);
        let max_i = row_index.max(col_index);

        // Record merge edge (R-style indexing)
        edges.push((group_mems[min_i], group_mems[max_i]));
        
        // Compute mllk at merge
        let merged = &sp_partition_counts[min_i] + &sp_partition_counts[max_i];
        let size_merged = partition_sizes[min_i] + partition_sizes[max_i];
        let mllk_merged = marginal_llk(
            &merged,
            size_merged,
            n_snps,
            &term1,
            &lg_table,
            &prior_index,
        );
        
        heights.push(mllk_merged);
        rk_vec.push(max_rk);

        // Update group membership (1-indexed internal node)
        group_mems[min_i] = (merge_step + 1) as i32;

        // Merge partitions
        sp_partition_counts[min_i] = merged;
        partition_sizes[min_i] = size_merged;
        
        // Update dk and ptree for merged partition
        let dk_new = log_sum_exp(
            lg_table.get(size_merged, 0),
            dk[min_i] + dk[max_i],
        );
        
        let ptree_new = log_sum_exp(
            lg_table.get(size_merged, 0) + mllk_merged - dk_new,
            dk[min_i] + dk[max_i] + ptree[min_i] + ptree[max_i] - dk_new,
        );
        
        dk[min_i] = dk_new;
        ptree[min_i] = ptree_new;

        // Mark max_i as used, invalidate its rk values
        used[max_i] = true;
        for i in 0..n_partitions {
            rk_matrix[max_i * n_partitions + i] = f64::NEG_INFINITY;
            rk_matrix[i * n_partitions + max_i] = f64::NEG_INFINITY;
        }

        // Recompute rk for merged partition against all remaining (parallel)
        let updates: Vec<(usize, f64)> = (0..n_partitions)
            .into_par_iter()
            .filter(|&pp| !used[pp] && pp != min_i)
            .map(|pp| {
                let merged = &sp_partition_counts[pp] + &sp_partition_counts[min_i];
                let size_merged = partition_sizes[pp] + partition_sizes[min_i];
                let mllk_val = marginal_llk(
                    &merged,
                    size_merged,
                    n_snps,
                    &term1,
                    &lg_table,
                    &prior_index,
                );
                let rk_val = mllk_val - ptree[min_i] - ptree[pp] 
                    + lg_table.get(size_merged, 0) - dk[min_i] - dk[pp];
                (pp, rk_val)
            })
            .collect();

        // Write back updates (symmetric)
        for (pp, rk_val) in updates {
            rk_matrix[min_i * n_partitions + pp] = rk_val;
            rk_matrix[pp * n_partitions + min_i] = rk_val;
        }
    }

    BhcResult {
        edges,
        heights,
        rk: rk_vec,
        initial_partition_llk,
    }
}

/// Build sparse count matrix for a single partition.
///
/// Reads SNP data in CSC format and accumulates allele counts across
/// the isolates in the partition.
///
/// # Arguments
/// * `partition` - List of 1-indexed isolate IDs
/// * `snp_data` - CSC sparse matrix data (1-4 for A/C/G/T)
/// * `snp_indices` - CSC row indices
/// * `snp_indptr` - CSC column pointers
/// * `n_snps` - Number of SNP positions
///
/// # Returns
/// Sparse count matrix (n_snps x 4) in CSR format
fn build_partition_counts(
    partition: &[usize],
    snp_data: &[u8],
    snp_indices: &[usize],
    snp_indptr: &[usize],
    n_snps: usize,
) -> CsMat<u32> {
    // Use TriMat for efficient construction
    let mut trimat = TriMat::new((n_snps, 4));

    for &isolate_id in partition {
        // Convert 1-indexed to 0-indexed
        let isolate_idx = isolate_id - 1;
        
        // Read CSC column for this isolate
        let start = snp_indptr[isolate_idx];
        let end = snp_indptr[isolate_idx + 1];
        
        for idx in start..end {
            let snp_idx = snp_indices[idx];
            let allele = snp_data[idx]; // 1-4 for A/C/G/T
            let allele_idx = (allele - 1) as usize; // Convert to 0-3
            
            // Accumulate count
            trimat.add_triplet(snp_idx, allele_idx, 1);
        }
    }

    trimat.to_csr()
}

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

    #[test]
    fn test_build_partition_counts() {
        // Create a simple CSC matrix (3 SNPs x 2 isolates)
        // Isolate 0: SNP 0 -> allele A (1), SNP 2 -> allele C (2)
        // Isolate 1: SNP 0 -> allele A (1), SNP 1 -> allele G (3)
        let snp_data = vec![1, 2, 1, 3];
        let snp_indices = vec![0, 2, 0, 1];
        let snp_indptr = vec![0, 2, 4];
        let n_snps = 3;
        
        // Partition with both isolates (1-indexed)
        let partition = vec![1, 2];
        
        let counts = build_partition_counts(
            &partition,
            &snp_data,
            &snp_indices,
            &snp_indptr,
            n_snps,
        );
        
        // Check dimensions
        assert_eq!(counts.rows(), 3);
        assert_eq!(counts.cols(), 4);
        
        // Check counts
        // SNP 0: 2x allele A (index 0)
        assert_eq!(counts.get(0, 0), Some(&2));
        // SNP 1: 1x allele G (index 2)
        assert_eq!(counts.get(1, 2), Some(&1));
        // SNP 2: 1x allele C (index 1)
        assert_eq!(counts.get(2, 1), Some(&1));
    }

    #[test]
    fn test_bhier_parallel_simple() {
        // Create minimal test case with 2 partitions
        // 4 isolates, 2 SNPs
        let snp_data = vec![1, 1, 2, 2]; // A, A, C, C
        let snp_indices = vec![0, 1, 0, 1];
        let snp_indptr = vec![0, 1, 2, 3, 4]; // 4 isolates
        let n_snps = 2;
        let n_isolates = 4;
        
        // Uniform prior
        let prior = Array2::from_elem((5, n_snps), 0.01);
        
        // Two partitions: [1,2] and [3,4]
        let partitions = vec![vec![1, 2], vec![3, 4]];
        
        // Initial dk (log of partition sizes)
        let dk_initial = vec![
            (2.0_f64).ln(),
            (2.0_f64).ln(),
        ];
        
        let result = bhier_parallel(
            &snp_data,
            &snp_indices,
            &snp_indptr,
            n_snps,
            n_isolates,
            &prior,
            &partitions,
            &dk_initial,
            1,
        );
        
        // Should have 1 merge (2 partitions -> 1)
        assert_eq!(result.edges.len(), 1);
        assert_eq!(result.heights.len(), 1);
        assert_eq!(result.rk.len(), 1);
        assert_eq!(result.initial_partition_llk.len(), 2);
        
        // Edge should be (-1, -2) in some order
        let edge = result.edges[0];
        assert!(edge == (-1, -2) || edge == (-2, -1));
        
        // Likelihoods should be finite
        assert!(result.heights[0].is_finite());
        assert!(result.rk[0].is_finite());
        assert!(result.initial_partition_llk[0].is_finite());
        assert!(result.initial_partition_llk[1].is_finite());
    }
}
