//! Multiple comparison correction methods
//!
//! Provides methods to control for multiple comparisons:
//! - FDR (False Discovery Rate) using Benjamini-Hochberg
//! - Bonferroni correction
//! - Holm-Bonferroni step-down procedure

/// Correction method for multiple comparisons
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CorrectionMethod {
    /// Benjamini-Hochberg FDR correction
    Fdr,
    /// Bonferroni correction (most conservative)
    Bonferroni,
    /// Holm-Bonferroni step-down procedure
    Holm,
    /// No correction
    None,
}

/// False Discovery Rate (FDR) correction using Benjamini-Hochberg procedure
///
/// Controls the expected proportion of false positives among rejected hypotheses.
///
/// # Arguments
/// * `pvalues` - Vector of uncorrected p-values
/// * `alpha` - Significance level (e.g., 0.05)
///
/// # Returns
/// Tuple of (reject, pvalues_corrected):
/// - reject: Vector of booleans indicating which hypotheses to reject
/// - pvalues_corrected: Adjusted p-values
///
/// # Example
/// ```
/// use rtx_neuro_stats::fdr_correction;
///
/// let pvalues = vec![0.01, 0.04, 0.03, 0.20, 0.001];
/// let (reject, pvals_corrected) = fdr_correction(&pvalues, 0.05);
///
/// // pvals_corrected[4] (0.001) should be smallest and still < 0.05
/// assert!(reject[4]); // 0.001 should be rejected
/// ```
pub fn fdr_correction(pvalues: &[f64], alpha: f64) -> (Vec<bool>, Vec<f64>) {
    let n = pvalues.len();
    if n == 0 {
        return (Vec::new(), Vec::new());
    }

    // Sort p-values with original indices
    let mut indexed: Vec<(usize, f64)> = pvalues.iter().copied().enumerate().collect();
    indexed.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));

    // Compute adjusted p-values using Benjamini-Hochberg
    let mut pvals_corrected = vec![0.0; n];
    let mut reject = vec![false; n];

    // Work backwards to ensure monotonicity
    let mut min_so_far: f64 = 1.0;
    for (rank, &(orig_idx, pval)) in indexed.iter().enumerate().rev() {
        // BH formula: p_adj = p * n / (rank + 1)
        let adjusted = (pval * n as f64 / (rank + 1) as f64).min(1.0);
        min_so_far = f64::min(min_so_far, adjusted);
        pvals_corrected[orig_idx] = min_so_far;
        reject[orig_idx] = min_so_far <= alpha;
    }

    (reject, pvals_corrected)
}

/// Bonferroni correction
///
/// Most conservative multiple comparison correction.
/// Simply multiplies p-values by the number of tests.
///
/// # Arguments
/// * `pvalues` - Vector of uncorrected p-values
/// * `alpha` - Significance level (e.g., 0.05)
///
/// # Returns
/// Tuple of (reject, pvalues_corrected)
pub fn bonferroni_correction(pvalues: &[f64], alpha: f64) -> (Vec<bool>, Vec<f64>) {
    let n = pvalues.len();
    if n == 0 {
        return (Vec::new(), Vec::new());
    }

    let pvals_corrected: Vec<f64> = pvalues.iter().map(|&p| (p * n as f64).min(1.0)).collect();

    let reject: Vec<bool> = pvals_corrected.iter().map(|&p| p <= alpha).collect();

    (reject, pvals_corrected)
}

/// Holm-Bonferroni step-down correction
///
/// Less conservative than Bonferroni while still controlling FWER.
///
/// # Arguments
/// * `pvalues` - Vector of uncorrected p-values
/// * `alpha` - Significance level (e.g., 0.05)
///
/// # Returns
/// Tuple of (reject, pvalues_corrected)
pub fn holm_correction(pvalues: &[f64], alpha: f64) -> (Vec<bool>, Vec<f64>) {
    let n = pvalues.len();
    if n == 0 {
        return (Vec::new(), Vec::new());
    }

    // Sort p-values with original indices
    let mut indexed: Vec<(usize, f64)> = pvalues.iter().copied().enumerate().collect();
    indexed.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));

    let mut pvals_corrected = vec![0.0; n];
    let mut reject = vec![false; n];

    // Holm step-down: p_adj[i] = max(p[j] * (n - j)) for j <= i
    let mut max_so_far: f64 = 0.0;
    for (rank, &(orig_idx, pval)) in indexed.iter().enumerate() {
        let adjusted = (pval * (n - rank) as f64).min(1.0);
        max_so_far = f64::max(max_so_far, adjusted);
        pvals_corrected[orig_idx] = max_so_far;
        reject[orig_idx] = max_so_far <= alpha;
    }

    (reject, pvals_corrected)
}

/// Apply any correction method
///
/// # Arguments
/// * `pvalues` - Vector of uncorrected p-values
/// * `alpha` - Significance level
/// * `method` - Correction method to use
///
/// # Returns
/// Tuple of (reject, pvalues_corrected)
pub fn correct_pvalues(
    pvalues: &[f64],
    alpha: f64,
    method: CorrectionMethod,
) -> (Vec<bool>, Vec<f64>) {
    match method {
        CorrectionMethod::Fdr => fdr_correction(pvalues, alpha),
        CorrectionMethod::Bonferroni => bonferroni_correction(pvalues, alpha),
        CorrectionMethod::Holm => holm_correction(pvalues, alpha),
        CorrectionMethod::None => {
            let reject: Vec<bool> = pvalues.iter().map(|&p| p <= alpha).collect();
            (reject, pvalues.to_vec())
        }
    }
}

/// Compute the number of significant tests after correction
pub fn n_significant(reject: &[bool]) -> usize {
    reject.iter().filter(|&&r| r).count()
}

/// Compute the proportion of significant tests
pub fn proportion_significant(reject: &[bool]) -> f64 {
    if reject.is_empty() {
        return 0.0;
    }
    n_significant(reject) as f64 / reject.len() as f64
}

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

    #[test]
    fn test_fdr_basic() {
        let pvalues = vec![0.01, 0.04, 0.03, 0.20, 0.001];
        let (reject, pvals_corrected) = fdr_correction(&pvalues, 0.05);

        // Check that 0.001 is still significant
        assert!(reject[4]);
        // 0.20 should not be significant
        assert!(!reject[3]);

        // Corrected p-values should be >= original
        for i in 0..pvalues.len() {
            assert!(pvals_corrected[i] >= pvalues[i]);
        }

        // Corrected p-values should be <= 1.0
        for &p in &pvals_corrected {
            assert!(p <= 1.0);
        }
    }

    #[test]
    fn test_fdr_all_significant() {
        let pvalues = vec![0.001, 0.002, 0.003, 0.004, 0.005];
        let (reject, _pvals_corrected) = fdr_correction(&pvalues, 0.05);

        // All should be significant
        assert!(reject.iter().all(|&r| r));
    }

    #[test]
    fn test_fdr_none_significant() {
        let pvalues = vec![0.5, 0.6, 0.7, 0.8, 0.9];
        let (reject, _pvals_corrected) = fdr_correction(&pvalues, 0.05);

        // None should be significant
        assert!(reject.iter().all(|&r| !r));
    }

    #[test]
    fn test_bonferroni() {
        let pvalues = vec![0.01, 0.02, 0.03];
        let (reject, pvals_corrected) = bonferroni_correction(&pvalues, 0.05);

        // 0.01 * 3 = 0.03, should be significant
        assert!(reject[0]);
        // 0.02 * 3 = 0.06, should not be significant
        assert!(!reject[1]);
        // 0.03 * 3 = 0.09, should not be significant
        assert!(!reject[2]);

        assert!((pvals_corrected[0] - 0.03).abs() < 1e-10);
        assert!((pvals_corrected[1] - 0.06).abs() < 1e-10);
        assert!((pvals_corrected[2] - 0.09).abs() < 1e-10);
    }

    #[test]
    fn test_holm() {
        let pvalues = vec![0.01, 0.04, 0.03, 0.20];
        let (reject, pvals_corrected) = holm_correction(&pvalues, 0.05);

        // Holm is less conservative than Bonferroni
        // 0.01 * 4 = 0.04, significant
        assert!(reject[0]);

        // Corrected values should be monotonic
        let mut sorted_corrected: Vec<(usize, f64)> = pvalues.iter().copied().enumerate().collect();
        sorted_corrected.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
    }

    #[test]
    fn test_correct_pvalues() {
        let pvalues = vec![0.01, 0.04, 0.10];

        let (rej_fdr, _) = correct_pvalues(&pvalues, 0.05, CorrectionMethod::Fdr);
        let (rej_bonf, _) = correct_pvalues(&pvalues, 0.05, CorrectionMethod::Bonferroni);
        let (rej_none, _) = correct_pvalues(&pvalues, 0.05, CorrectionMethod::None);

        // No correction should have more rejections
        assert!(n_significant(&rej_none) >= n_significant(&rej_fdr));
        assert!(n_significant(&rej_fdr) >= n_significant(&rej_bonf));
    }

    #[test]
    fn test_empty_input() {
        let pvalues: Vec<f64> = Vec::new();

        let (reject, pvals_corrected) = fdr_correction(&pvalues, 0.05);
        assert!(reject.is_empty());
        assert!(pvals_corrected.is_empty());
    }
}