// GPL-3.0-or-later
// Copyright 2024--present rsx-rs developers

//! Statistical functions: chi-squared test with Yates correction,
//! Bonferroni multiple testing correction, and group bias.

use std::fmt;

/// Format a float like C++ operator<< default: %g with 6 significant digits.
/// This matches the C++ radsex output format exactly.
pub struct Cg(pub f64);

impl fmt::Display for Cg {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // %g with 6 significant digits, strip trailing zeros.
        // Exponent computed safely to avoid log10() floor off-by-one on
        // powers of 10 and values near them (common fp hazard).
        let val = self.0;
        if val == 0.0 {
            return write!(f, "0");
        }
        let abs = val.abs();
        let exp = safe_float_exponent(abs);
        if (-4..6).contains(&exp) {
            // Fixed notation
            let precision = (5 - exp).max(0) as usize;
            let formatted = format!("{:.*}", precision, val);
            if formatted.contains('.') {
                let trimmed = formatted.trim_end_matches('0').trim_end_matches('.');
                write!(f, "{trimmed}")
            } else {
                write!(f, "{formatted}")
            }
        } else {
            // Scientific notation
            let mantissa = val / 10f64.powi(exp);
            let formatted = format!("{:.5}", mantissa);
            let trimmed = formatted.trim_end_matches('0').trim_end_matches('.');
            write!(f, "{trimmed}e{exp:+03}")
        }
    }
}

/// Robust base-10 exponent for %g-style formatting.
/// Uses log10 + explicit boundary check/adjustment to handle
/// cases where f64 log10(10^k) != k exactly.
#[inline]
fn safe_float_exponent(abs: f64) -> i32 {
    if abs == 0.0 || !abs.is_finite() {
        return 0;
    }
    let mut e = abs.log10().floor() as i32;
    // Verify 10^e <= abs < 10^(e+1), adjust for fp error
    if 10f64.powi(e) > abs && e > i32::MIN {
        e -= 1;
    }
    if 10f64.powi(e + 1) <= abs {
        e += 1;
    }
    e
}

/// Chi-squared statistic with Yates continuity correction for a 2x2 table.
///
/// Implements the shortcut formula::
///
///   chi2 = N * (|ad - bc| - N/2)^2 / (a+b)(c+d)(a+c)(b+d)
///
/// where the contingency table is::
///
///   |           | marker present | marker absent |
///   |-----------|----------------|---------------|
///   | group1    | n_group1       | total1 - n1   |
///   | group2    | n_group2       | total2 - n2   |
pub fn chi_squared_yates(
    n_group1: u32,
    n_group2: u32,
    total_group1: u32,
    total_group2: u32,
) -> f64 {
    // Use u64 totals so large valid u32 inputs do not overflow during addition.
    let n = (total_group1 as u64 + total_group2 as u64) as f64;
    let ns = total_group1 as f64;
    let nf = total_group2 as f64;
    let na = (n_group1 as u64 + n_group2 as u64) as f64;
    let nb = n - na;

    // u32 cross products fit within u64 across the full input domain.
    let ad = (n_group1 as u64) * (total_group2 as u64);
    let bc = (n_group2 as u64) * (total_group1 as u64);
    let ad_bc = ad.abs_diff(bc) as f64;
    let yates = (ad_bc - n / 2.0).max(0.0);

    n * yates * yates / (ns * nf * na * nb)
}

/// P-value for a chi-squared statistic with df=1.
///
/// Uses the exact identity: for df=1, the chi-squared CDF is
///   P(chi2) = erf(sqrt(chi2/2))
/// so the p-value is:
///   p = 1 - P(chi2) = erfc(sqrt(chi2/2))
///
/// This replaces the full regularized gamma function with a single
/// libm erfc call. Derived via SymPy::
///
///   gamma(1/2, x) = sqrt(pi) * erf(sqrt(x))
///   Gamma(1/2) = sqrt(pi)
///   P(1/2, x) = erf(sqrt(x))
///   p = erfc(sqrt(chi2/2))
pub fn chi_squared_p(chi_sq: f64) -> f64 {
    if chi_sq.is_nan() || chi_sq <= 0.0 {
        return 1.0;
    }
    fast_erfc((chi_sq / 2.0).sqrt()).min(1.0)
}

/// erfc(t) via libm (CPU-optimized, ~30ns).
/// For GPU kernels, use `fast_erfc_poly` which avoids libm calls.
#[inline]
fn fast_erfc(t: f64) -> f64 {
    libm::erfc(t)
}

/// Sollya-generated minimax polynomial for erfc(t) on [0, 6].
/// Single degree-40 polynomial -- branchless, no exp(), GPU/SIMD ready.
/// Max absolute error: 8.2e-17 (below f64 epsilon).
/// See `scripts/sollya/erfc_direct.sollya`.
#[inline]
#[allow(dead_code)]
pub fn fast_erfc_poly(t: f64) -> f64 {
    if t >= 6.0 {
        return 0.0;
    }
    horner(t, &ERFC_POLY)
}

/// Horner evaluation: P(x) = c[0] + x*(c[1] + x*(c[2] + ...))
/// Uses FMA (mul_add) for maximum precision and speed.
#[inline(always)]
fn horner(x: f64, coeffs: &[f64]) -> f64 {
    let mut result = 0.0f64;
    for &c in coeffs.iter().rev() {
        result = result.mul_add(x, c);
    }
    result
}

// Sollya remez degree-40 minimax polynomial for erfc(t) on [0, 6].
// Max abs error: 8.22e-17. No exp() call. SIMD/GPU friendly.
// Coefficients as hex float via hexf! -- direct Sollya output, zero transcription error.
// Generated by: cd scripts && pixi run minimax-direct
#[rustfmt::skip]
const ERFC_POLY: [f64; 41] = [
    hexf::hexf64!("0x1.fffffffffffffp-1"),   // c0
    hexf::hexf64!("-0x1.20dd750429abbp0"),    // c1
    hexf::hexf64!("-0x1.b83ea9b36d55dp-39"),  // c2
    hexf::hexf64!("0x1.812746b1e8a4cp-2"),    // c3
    hexf::hexf64!("-0x1.c01f0b102e8eep-30"),  // c4
    hexf::hexf64!("-0x1.ce2f1d3ee338ep-4"),   // c5
    hexf::hexf64!("-0x1.c92fce83a23c7p-24"),  // c6
    hexf::hexf64!("0x1.b82ed8ff9fbe4p-6"),    // c7
    hexf::hexf64!("-0x1.6af32fff4be83p-20"),  // c8
    hexf::hexf64!("-0x1.56386821f01b1p-8"),   // c9
    hexf::hexf64!("0x1.84dc6b9cac39ap-20"),   // c10
    hexf::hexf64!("0x1.b5ac4763df8f1p-11"),   // c11
    hexf::hexf64!("0x1.226e298aa6414p-14"),   // c12
    hexf::hexf64!("-0x1.24f7e6a2b2ceap-12"),  // c13
    hexf::hexf64!("0x1.239ab445441dcp-12"),   // c14
    hexf::hexf64!("-0x1.8b31805e161c4p-12"),  // c15
    hexf::hexf64!("0x1.dee4be601e3e9p-12"),   // c16
    hexf::hexf64!("-0x1.d77acde978739p-12"),  // c17
    hexf::hexf64!("0x1.883aaa14debap-12"),    // c18
    hexf::hexf64!("-0x1.18dadbd5cc3abp-12"),  // c19
    hexf::hexf64!("0x1.5b8d481772fb4p-13"),   // c20
    hexf::hexf64!("-0x1.74855604ef48ap-14"),  // c21
    hexf::hexf64!("0x1.5acc93d93d3b4p-15"),   // c22
    hexf::hexf64!("-0x1.18f50a2c68becp-16"),  // c23
    hexf::hexf64!("0x1.8c5fd35d0866cp-18"),   // c24
    hexf::hexf64!("-0x1.e6a1ce3798ffap-20"),  // c25
    hexf::hexf64!("0x1.03a13e31b65dp-21"),    // c26
    hexf::hexf64!("-0x1.e0be7f87c8eaep-24"),  // c27
    hexf::hexf64!("0x1.815459515f217p-26"),   // c28
    hexf::hexf64!("-0x1.0a9d538fdfb12p-28"),  // c29
    hexf::hexf64!("0x1.3d4235505e0e1p-31"),   // c30
    hexf::hexf64!("-0x1.42ee4b10402d8p-34"),  // c31
    hexf::hexf64!("0x1.172dd0bdac958p-37"),   // c32
    hexf::hexf64!("-0x1.960c57cfddb97p-41"),  // c33
    hexf::hexf64!("0x1.ea4beaf5f9701p-45"),   // c34
    hexf::hexf64!("-0x1.e2a4b9cf888a5p-49"),  // c35
    hexf::hexf64!("0x1.796af39619721p-53"),   // c36
    hexf::hexf64!("-0x1.c31352aaa9451p-58"),  // c37
    hexf::hexf64!("0x1.82f8c94463357p-63"),   // c38
    hexf::hexf64!("-0x1.a861a25d26dp-69"),    // c39
    hexf::hexf64!("0x1.bef061298bc36p-76"),   // c40
];

/// Compute p-value of association with group using chi-squared test
/// with Yates correction. Matches C++ `get_p_association` exactly.
pub fn p_association(n_group1: u32, n_group2: u32, total_group1: u32, total_group2: u32) -> f64 {
    let chi_sq = chi_squared_yates(n_group1, n_group2, total_group1, total_group2);

    // NaN guard: if chi_sq is NaN (e.g. division by zero when marker
    // is present in all individuals), return p = 1.0
    let p = if chi_sq.is_nan() {
        1.0
    } else {
        chi_squared_p(chi_sq).min(1.0)
    };

    // Lower bound: 1e-16 (matching C++ behavior)
    p.max(1e-16)
}

/// Apply Bonferroni correction to a p-value.
pub fn bonferroni_correct(p: f64, n_markers: u64) -> f64 {
    (p * n_markers as f64).min(1.0)
}

/// Group bias: difference in marker frequency between two groups.
/// Ranges from -1.0 (only in group2) to +1.0 (only in group1).
pub fn group_bias(n_group1: u32, total_group1: u32, n_group2: u32, total_group2: u32) -> f64 {
    (n_group1 as f64 / total_group1 as f64) - (n_group2 as f64 / total_group2 as f64)
}

/// Find median of a mutable slice (partially reorders in-place).
///
/// Uses order-statistic selection (`select_nth_unstable`) at index `len/2`
/// instead of a full sort. This is the **upper median** for even length:
/// after zeros, the element at conceptual position `len/2` (0-based), matching
/// the streaming external-sort path. Selection is O(n) average vs O(n log n)
/// sort; proven equivalent to sorting then indexing `len/2` (see
/// `proofs/lean/MedianSelect.lean` and `scripts/sympy/median_select_proof.py`).
pub fn find_median(data: &mut [u16]) -> u16 {
    let len = data.len();
    if len == 0 {
        return 0;
    }
    let k = len / 2;
    // Partition so data[k] is the element that would sit at index k if fully sorted.
    let (_left, mid, _right) = data.select_nth_unstable(k);
    *mid
}

// ========================================================================
// Alternative statistical tests
// ========================================================================

/// G-test (log-likelihood ratio) for 2x2 table.
/// Better asymptotic properties than chi-squared.
pub fn g_test(n_g1: u32, n_g2: u32, total_g1: u32, total_g2: u32) -> f64 {
    let a = n_g1 as f64;
    let b = (total_g1 - n_g1) as f64;
    let c = n_g2 as f64;
    let d = (total_g2 - n_g2) as f64;
    let n = (total_g1 + total_g2) as f64;

    let mut g = 0.0f64;
    let expected = |_obs: f64, col_total: f64, row_total: f64| -> f64 { row_total * col_total / n };

    let row1 = a + c;
    let row2 = b + d;
    let col1 = a + b;
    let col2 = c + d;

    for &(obs, rt, ct) in &[
        (a, row1, col1),
        (b, row2, col1),
        (c, row1, col2),
        (d, row2, col2),
    ] {
        if obs > 0.0 {
            let exp = expected(obs, ct, rt);
            if exp > 0.0 {
                g += obs * (obs / exp).ln();
            }
        }
    }
    g *= 2.0;

    // P-value: same chi-squared CDF with df=1
    if g.is_nan() || g <= 0.0 {
        1.0
    } else {
        chi_squared_p(g).clamp(1e-16, 1.0)
    }
}

/// Fisher's exact test for 2x2 table (one-sided, more extreme).
/// Uses hypergeometric probability. Better than chi-squared for small n.
pub fn fisher_exact(n_g1: u32, n_g2: u32, total_g1: u32, total_g2: u32) -> f64 {
    let a = n_g1;
    let b = total_g1 - n_g1;
    let c = n_g2;
    let d = total_g2 - n_g2;
    let n = total_g1 + total_g2;
    let row1 = a + c;
    let col1 = a + b;

    // P(table) = C(row1,a) * C(row2,b) / C(n, col1)
    // Sum probabilities of all tables as or more extreme
    let log_p_observed = log_hypergeometric(a, b, c, d, n, row1, col1);

    let min_a = col1.saturating_sub(n - row1);
    let max_a = row1.min(col1);

    // Online log-sum-exp over the density tail (tables as or more extreme).
    // Same math as collecting a Vec then logsumexp, without per-call heap alloc
    // (signif --test fisher invokes this once per marker).
    // Include if log_p <= observed (within tiny fp tolerance for lgamma error).
    let mut max_log = f64::NEG_INFINITY;
    let mut sum_rel = 0.0f64;
    let mut any = false;
    for a_i in min_a..=max_a {
        let b_i = col1 - a_i;
        let c_i = row1 - a_i;
        let d_i = (n - row1) - b_i;
        let log_p = log_hypergeometric(a_i, b_i, c_i, d_i, n, row1, col1);
        if log_p <= log_p_observed + 1e-9 {
            // Standard online logsumexp update (Goldberg / Press et al.)
            if !any {
                max_log = log_p;
                sum_rel = 1.0;
                any = true;
            } else if log_p > max_log {
                sum_rel = sum_rel * (max_log - log_p).exp() + 1.0;
                max_log = log_p;
            } else {
                sum_rel += (log_p - max_log).exp();
            }
        }
    }

    let p_sum = if !any {
        0.0
    } else {
        (max_log + sum_rel.ln()).exp()
    };

    p_sum.clamp(1e-16, 1.0)
}

fn log_hypergeometric(a: u32, b: u32, c: u32, d: u32, n: u32, row1: u32, col1: u32) -> f64 {
    let col2 = n - col1;
    let row2 = n - row1;
    lfact(row1) + lfact(row2) + lfact(col1) + lfact(col2)
        - lfact(n)
        - lfact(a)
        - lfact(b)
        - lfact(c)
        - lfact(d)
}

fn lfact(n: u32) -> f64 {
    libm::lgamma((n + 1) as f64)
}

// ========================================================================
// Benjamini-Hochberg FDR
// ========================================================================

/// Apply Benjamini-Hochberg FDR correction to a vector of p-values.
/// Returns adjusted p-values (q-values). Controls FDR at the given level.
pub fn benjamini_hochberg(p_values: &[f64]) -> Vec<f64> {
    let m = p_values.len();
    if m == 0 {
        return Vec::new();
    }

    // Sort indices by p-value
    let mut indices: Vec<usize> = (0..m).collect();
    indices.sort_by(|&a, &b| p_values[a].total_cmp(&p_values[b]));

    let mut q_values = vec![0.0f64; m];
    let mut cummin = f64::INFINITY;

    for (rank_rev, &idx) in indices.iter().rev().enumerate() {
        let rank = m - rank_rev; // 1-based rank from sorted order
        let adjusted = p_values[idx] * m as f64 / rank as f64;
        cummin = cummin.min(adjusted);
        q_values[idx] = cummin.min(1.0);
    }

    q_values
}

// ========================================================================
// Bayesian methods
// ========================================================================

/// Bayes Factor for association in a 2x2 contingency table.
/// BF > 1: evidence for association. BF > 10: strong evidence.
/// Uses Beta-Binomial marginal likelihoods with uniform Beta(1,1) priors.
/// H1: separate proportions per group. H0: shared proportion.
pub fn bayes_factor_2x2(n_g1: u32, n_g2: u32, total_g1: u32, total_g2: u32) -> f64 {
    // H1 (association): p_g1, p_g2 independent, each Beta(1,1)
    // Marginal = B(k+1, n-k+1) / B(1,1) for each group
    let log_h1 = log_beta_binom(n_g1, total_g1) + log_beta_binom(n_g2, total_g2);

    // H0 (independence): single shared p ~ Beta(1,1)
    let log_h0 = log_beta_binom(n_g1 + n_g2, total_g1 + total_g2);

    (log_h1 - log_h0).exp()
}

/// Log marginal likelihood for Binomial(k|n) with Beta(1,1) prior.
fn log_beta_binom(k: u32, n: u32) -> f64 {
    libm::lgamma((k + 1) as f64) + libm::lgamma((n - k + 1) as f64) - libm::lgamma((n + 2) as f64)
}

/// Posterior probability that a marker is sex-linked (empirical Bayes).
/// Uses a Beta-Binomial conjugate model.
/// pi: prior probability of sex-linkage (estimated from data or set to 0.01).
/// p_sex: assumed frequency in the linked sex (e.g., 0.9).
pub fn posterior_sex_linked(
    n_g1: u32,
    n_g2: u32,
    total_g1: u32,
    total_g2: u32,
    pi: f64,
    p_sex: f64,
) -> f64 {
    let ll_g1_linked =
        binom_logpmf(n_g1, total_g1, p_sex) + binom_logpmf(n_g2, total_g2, 1.0 - p_sex);
    let ll_g2_linked =
        binom_logpmf(n_g1, total_g1, 1.0 - p_sex) + binom_logpmf(n_g2, total_g2, p_sex);
    let ll_linked = logsumexp2(ll_g1_linked, ll_g2_linked) - 2.0f64.ln();
    let ll_null = binom_logpmf(n_g1, total_g1, 0.5) + binom_logpmf(n_g2, total_g2, 0.5);

    let log_odds = ll_linked - ll_null + (pi / (1.0 - pi)).ln();

    // Posterior = 1 / (1 + exp(-log_odds))
    if log_odds > 20.0 {
        1.0 // overflow guard
    } else if log_odds < -20.0 {
        0.0
    } else {
        1.0 / (1.0 + (-log_odds).exp())
    }
}

fn logsumexp2(a: f64, b: f64) -> f64 {
    let max = a.max(b);
    if max.is_infinite() {
        max
    } else {
        max + ((a - max).exp() + (b - max).exp()).ln()
    }
}

fn binom_logpmf(k: u32, n: u32, p: f64) -> f64 {
    if p <= 0.0 || p >= 1.0 {
        if (p <= 0.0 && k == 0) || (p >= 1.0 && k == n) {
            return 0.0;
        }
        return f64::NEG_INFINITY;
    }
    lfact(n) - lfact(k) - lfact(n - k) + k as f64 * p.ln() + (n - k) as f64 * (1.0 - p).ln()
}

/// Empirical Bayes EM: estimate pi (fraction of sex-linked markers) from data.
/// Returns (pi, posteriors) after convergence.
/// group_counts: Vec of (n_g1, n_g2) for each marker.
pub fn empirical_bayes_em(
    group_counts: &[(u32, u32)],
    total_g1: u32,
    total_g2: u32,
    p_sex: f64,
    max_iter: usize,
) -> (f64, Vec<f64>) {
    let m = group_counts.len();
    let mut pi = 0.01f64; // initial estimate
    let mut posteriors = vec![0.0f64; m];

    for _iter in 0..max_iter {
        // E-step: compute posteriors
        for (i, &(n_g1, n_g2)) in group_counts.iter().enumerate() {
            posteriors[i] = posterior_sex_linked(n_g1, n_g2, total_g1, total_g2, pi, p_sex);
        }

        // M-step: update pi
        let new_pi: f64 = posteriors.iter().sum::<f64>() / m as f64;
        if (new_pi - pi).abs() < 1e-8 {
            pi = new_pi;
            break;
        }
        pi = new_pi;
    }

    // Final E-step with converged pi
    for (i, &(n_g1, n_g2)) in group_counts.iter().enumerate() {
        posteriors[i] = posterior_sex_linked(n_g1, n_g2, total_g1, total_g2, pi, p_sex);
    }

    (pi, posteriors)
}

// ========================================================================
// Logistic regression (population-aware, Newton-Raphson)
// ========================================================================

/// Logistic regression: fit y ~ X using IRLS (Newton-Raphson).
/// X: n x p design matrix (row-major), y: n binary outcomes (0/1).
/// Returns coefficient vector beta (length p).
pub fn logistic_regression(x: &[f64], y: &[f64], n: usize, p: usize, max_iter: usize) -> Vec<f64> {
    let mut beta = vec![0.0f64; p];

    for _iter in 0..max_iter {
        // Compute predictions: mu = sigmoid(X @ beta)
        let mut mu = vec![0.0f64; n];
        for i in 0..n {
            let mut eta = 0.0;
            for j in 0..p {
                eta += x[i * p + j] * beta[j];
            }
            mu[i] = 1.0 / (1.0 + (-eta).exp());
            // Clamp to avoid log(0)
            mu[i] = mu[i].clamp(1e-10, 1.0 - 1e-10);
        }

        // Gradient: X^T @ (y - mu)
        let mut grad = vec![0.0f64; p];
        for j in 0..p {
            for i in 0..n {
                grad[j] += x[i * p + j] * (y[i] - mu[i]);
            }
        }

        // Hessian: -X^T @ W @ X where W = diag(mu * (1 - mu))
        let mut hessian = vec![0.0f64; p * p];
        for i in 0..n {
            let w = mu[i] * (1.0 - mu[i]);
            for j1 in 0..p {
                for j2 in 0..p {
                    hessian[j1 * p + j2] -= x[i * p + j1] * w * x[i * p + j2];
                }
            }
        }

        // Solve: delta = -H^{-1} @ grad (using Gaussian elimination for small p)
        let delta = solve_linear(&hessian, &grad, p);

        // Update beta
        let mut max_delta = 0.0f64;
        for j in 0..p {
            beta[j] -= delta[j];
            max_delta = max_delta.max(delta[j].abs());
        }

        if max_delta < 1e-8 {
            break;
        }
    }

    beta
}

/// Solve Ax = b via Gaussian elimination (for small p x p systems).
fn solve_linear(a: &[f64], b: &[f64], n: usize) -> Vec<f64> {
    let mut aug = vec![0.0f64; n * (n + 1)];
    for i in 0..n {
        for j in 0..n {
            aug[i * (n + 1) + j] = a[i * n + j];
        }
        aug[i * (n + 1) + n] = b[i];
    }

    // Forward elimination
    for col in 0..n {
        // Partial pivoting
        let mut max_row = col;
        for row in (col + 1)..n {
            if aug[row * (n + 1) + col].abs() > aug[max_row * (n + 1) + col].abs() {
                max_row = row;
            }
        }
        if max_row != col {
            for j in 0..=n {
                aug.swap(col * (n + 1) + j, max_row * (n + 1) + j);
            }
        }

        let pivot = aug[col * (n + 1) + col];
        if pivot.abs() < 1e-15 {
            continue;
        }

        for row in (col + 1)..n {
            let factor = aug[row * (n + 1) + col] / pivot;
            for j in col..=n {
                aug[row * (n + 1) + j] -= factor * aug[col * (n + 1) + j];
            }
        }
    }

    // Back substitution
    let mut x = vec![0.0f64; n];
    for i in (0..n).rev() {
        let pivot = aug[i * (n + 1) + i];
        if pivot.abs() < 1e-15 {
            continue;
        }
        x[i] = aug[i * (n + 1) + n];
        for j in (i + 1)..n {
            x[i] -= aug[i * (n + 1) + j] * x[j];
        }
        x[i] /= pivot;
    }
    x
}

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

    #[test]
    fn test_chi_squared_known_values() {
        // 10 males with marker, 0 females, out of 15 males and 10 females
        let chi_sq = chi_squared_yates(10, 0, 15, 10);
        assert!(chi_sq > 0.0);
        let p = chi_squared_p(chi_sq);
        assert!(p < 0.05);
    }

    #[test]
    fn test_equal_distribution_not_significant() {
        // 5 males, 5 females, out of 10 each -- no association
        let p = p_association(5, 5, 10, 10);
        assert!(p > 0.05);
    }

    #[test]
    fn test_p_association_floor() {
        // Extreme case: should be capped at 1e-16
        let p = p_association(20, 0, 20, 20);
        assert!(p >= 1e-16);
    }

    #[test]
    fn test_chi_squared_yates_large_values_no_overflow() {
        // These inputs make the u32 cross product exceed i64::MAX while fitting in u64.
        let ng1: u32 = 3_000_000_000;
        let ng2: u32 = 3_000_000_000;
        let tg1: u32 = 4_000_000_000;
        let tg2: u32 = 4_000_000_000;
        let chi = chi_squared_yates(ng1, ng2, tg1, tg2);
        assert!(chi.is_finite(), "chi2 must be finite, got {chi}");
        // Cross-check with explicit wide arithmetic (u64 intermediates)
        let n = (tg1 as u64 + tg2 as u64) as f64;
        let ad = (ng1 as u64) * (tg2 as u64);
        let bc = (ng2 as u64) * (tg1 as u64);
        let adbc = ad.abs_diff(bc) as f64;
        let y = (adbc - n / 2.0).max(0.0);
        let na = (ng1 as u64 + ng2 as u64) as f64;
        let nb = n - na;
        let ns = tg1 as f64;
        let nf = tg2 as f64;
        let expected = n * y * y / (ns * nf * na * nb);
        assert!(
            (chi - expected).abs() < 1e-4,
            "chi2 mismatch: {chi} vs {expected}"
        );
    }

    #[test]
    fn test_bonferroni() {
        assert_eq!(bonferroni_correct(0.01, 10), 0.1);
        assert_eq!(bonferroni_correct(0.5, 10), 1.0); // capped at 1.0
    }

    #[test]
    fn test_group_bias() {
        let bias = group_bias(10, 10, 0, 10);
        assert!((bias - 1.0).abs() < f64::EPSILON);
        let bias = group_bias(0, 10, 10, 10);
        assert!((bias - (-1.0)).abs() < f64::EPSILON);
    }

    #[test]
    fn test_find_median_odd() {
        let mut data = vec![3, 1, 2];
        assert_eq!(find_median(&mut data), 2);
    }

    #[test]
    fn test_find_median_even() {
        let mut data = vec![4, 1, 3, 2];
        // Upper median for even length (data[len/2] after sort): unifies with
        // streaming external-sort median in depth command (no fractional depths).
        assert_eq!(find_median(&mut data), 3); // sorted [1,2,3,4], position 2 -> 3
    }

    /// Selection at k=len/2 must match full sort then index k (upper median).
    #[test]
    fn test_find_median_matches_sorted_index() {
        for seed in 0u64..40 {
            let mut rng = seed;
            let n = 1 + (rng % 64) as usize;
            let mut data: Vec<u16> = (0..n)
                .map(|_| {
                    rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1);
                    (rng >> 33) as u16
                })
                .collect();
            let mut sorted = data.clone();
            sorted.sort_unstable();
            let expected = sorted[n / 2];
            assert_eq!(find_median(&mut data), expected, "seed={seed} n={n}");
        }
    }

    // === G-test ===
    #[test]
    fn test_g_test_no_association() {
        let p = g_test(5, 5, 10, 10);
        assert!(
            p > 0.05,
            "equal distribution should not be significant: p={p}"
        );
    }

    #[test]
    fn test_g_test_strong_association() {
        let p = g_test(10, 0, 10, 10);
        assert!(
            p < 0.01,
            "all-male marker should be highly significant: p={p}"
        );
    }

    // === Fisher's exact ===
    #[test]
    fn test_fisher_exact_no_association() {
        let p = fisher_exact(3, 3, 6, 6);
        assert!(p > 0.05, "equal distribution: p={p}");
    }

    #[test]
    fn test_fisher_exact_strong_association() {
        let p = fisher_exact(6, 0, 6, 6);
        assert!(p < 0.05, "all-male marker: p={p}");
    }

    #[test]
    fn test_fisher_exact_small_sample() {
        // Fisher's should work well even with tiny samples
        let p = fisher_exact(2, 0, 2, 2);
        assert!(p > 0.0 && p <= 1.0, "valid p-value: p={p}");
    }

    #[test]
    fn test_fisher_exact_strong_signal_stable_sum() {
        // Strong association with enough tables that naive exp() sum would underflow many terms.
        // With logsumexp the p should still be correctly tiny but > 1e-16.
        let p = fisher_exact(20, 0, 25, 25);
        assert!(
            p > 0.0 && p < 1e-6,
            "strong signal should give very small p, got {p}"
        );
        assert!(
            p >= 1e-16,
            "p should not underflow to the floor incorrectly"
        );
    }

    // === Benjamini-Hochberg ===
    #[test]
    fn test_bh_fdr_basic() {
        let pvals = vec![0.01, 0.04, 0.03, 0.20, 0.50];
        let q = benjamini_hochberg(&pvals);
        assert_eq!(q.len(), 5);
        // All q-values should be >= corresponding p-values
        for (p, q) in pvals.iter().zip(q.iter()) {
            assert!(*q >= *p, "q={q} < p={p}");
        }
        // q-values should be <= 1.0
        for &qi in &q {
            assert!(qi <= 1.0);
        }
    }

    #[test]
    fn test_bh_fdr_monotone() {
        let pvals = vec![0.001, 0.01, 0.05, 0.1, 0.5];
        let q = benjamini_hochberg(&pvals);
        // Sorted q-values should be non-decreasing
        let mut sorted_q = q.clone();
        sorted_q.sort_by(|a, b| a.total_cmp(b));
        // The adjusted values for sorted p-values should be non-decreasing
    }

    // === Bayes Factor ===
    #[test]
    fn test_bayes_factor_no_association() {
        let bf = bayes_factor_2x2(5, 5, 10, 10);
        assert!(bf < 3.0, "equal distribution BF should be < 3: BF={bf}");
    }

    #[test]
    fn test_bayes_factor_strong_association() {
        let bf = bayes_factor_2x2(10, 0, 10, 10);
        assert!(bf > 10.0, "all-male marker BF should be > 10: BF={bf}");
    }

    // === Posterior sex-linked ===
    #[test]
    fn test_posterior_strong_signal() {
        let post = posterior_sex_linked(10, 0, 10, 10, 0.01, 0.9);
        assert!(
            post > 0.9,
            "strong signal should have high posterior: {post}"
        );
    }

    #[test]
    fn test_posterior_strong_signal_is_symmetric() {
        let group1 = posterior_sex_linked(10, 0, 10, 10, 0.01, 0.9);
        let group2 = posterior_sex_linked(0, 10, 10, 10, 0.01, 0.9);
        assert!(
            group1 > 0.9,
            "group1-linked posterior should be high: {group1}"
        );
        assert!(
            group2 > 0.9,
            "group2-linked posterior should be high: {group2}"
        );
        assert!(
            (group1 - group2).abs() < 1e-12,
            "posterior should be symmetric: group1={group1}, group2={group2}"
        );
    }

    #[test]
    fn test_posterior_no_signal() {
        let post = posterior_sex_linked(5, 5, 10, 10, 0.01, 0.9);
        assert!(
            post < 0.1,
            "equal distribution should have low posterior: {post}"
        );
    }

    // === Empirical Bayes EM ===
    #[test]
    fn test_em_convergence() {
        let counts = vec![
            (10, 0),
            (9, 1),
            (10, 0), // sex-linked
            (5, 5),
            (4, 6),
            (6, 4),
            (5, 5),
            (5, 5),
            (5, 5),
            (5, 5), // null
        ];
        let (pi, posteriors) = empirical_bayes_em(&counts, 10, 10, 0.9, 50);
        assert!(pi > 0.0 && pi < 1.0, "pi should converge: {pi}");
        // Sex-linked markers should have higher posteriors
        let avg_linked: f64 = posteriors[..3].iter().sum::<f64>() / 3.0;
        let avg_null: f64 = posteriors[3..].iter().sum::<f64>() / 7.0;
        assert!(
            avg_linked > avg_null,
            "linked={avg_linked} should > null={avg_null}"
        );
    }

    // === Logistic regression ===
    #[test]
    fn test_logistic_basic() {
        // Simple separable data: 4 samples, 1 predictor
        let x = vec![
            1.0, -2.0, // intercept, x for sample 0
            1.0, -1.0, 1.0, 1.0, 1.0, 2.0,
        ];
        let y = vec![0.0, 0.0, 1.0, 1.0];
        let beta = logistic_regression(&x, &y, 4, 2, 50);
        assert_eq!(beta.len(), 2);
        // beta[1] should be positive (higher x -> higher P(y=1))
        assert!(beta[1] > 0.0, "slope should be positive: {}", beta[1]);
    }

    // === Cg formatter (must match C++ %g 6 sigdig + trim for golden files) ===
    #[test]
    fn test_cg_basic_and_small_p() {
        // Common p-value outputs from chi_squared_p
        assert_eq!(format!("{}", Cg(0.05)), "0.05");
        assert_eq!(format!("{}", Cg(0.01)), "0.01");
        assert_eq!(format!("{}", Cg(1e-16)), "1e-16");
        assert_eq!(format!("{}", Cg(0.0)), "0");
        assert_eq!(format!("{}", Cg(1.0)), "1");
    }

    #[test]
    fn test_cg_significant_digits_and_scientific() {
        // 6 significant digits, trailing zeros trimmed, switches to sci for |exp| large
        assert_eq!(format!("{}", Cg(1.23456789)), "1.23457");
        assert_eq!(format!("{}", Cg(123456.789)), "123457");
        assert_eq!(format!("{}", Cg(1.23456e-10)), "1.23456e-10");
        assert_eq!(format!("{}", Cg(9.99999e-5)), "9.99999e-05");
    }

    #[test]
    fn test_cg_used_in_p_association_matches_expected_strings() {
        // End-to-end: p values that appear in distrib/signif outputs must format consistently
        let p = p_association(10, 0, 15, 10);
        let s = format!("{}", Cg(p));
        // Must be a short non-scientific or e- form, never garbage
        assert!(s.len() <= 12 && (s.contains('.') || s.contains('e')));
    }

    #[test]
    fn test_cg_exponent_boundary_cases() {
        // These values are known to cause log10().floor() off-by-one in naive impls
        assert_eq!(format!("{}", Cg(10.0)), "10");
        assert_eq!(format!("{}", Cg(100.0)), "100");
        assert_eq!(format!("{}", Cg(0.1)), "0.1");
        assert_eq!(format!("{}", Cg(0.01)), "0.01");
        assert_eq!(format!("{}", Cg(1e-4)), "0.0001");
        // Very close to power of 10 from below
        assert_eq!(format!("{}", Cg(9.999999999999998)), "10");
        // Small scientific that must not become 0.000099999...
        let s = format!("{}", Cg(9.99999e-5));
        assert!(s == "9.99999e-05" || s == "0.0001" || s.starts_with("1e-04"));
    }
}
