regression-diagnostics 0.2.0

use statrs::distribution::{ChiSquared, ContinuousCDF};

use crate::OlsFit;

/// Result of the Jarque-Bera normality test.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct JarqueBera {
    /// The JB test statistic.
    pub statistic: f64,
    /// Upper-tail p-value under the χ²(2) distribution.
    pub p_value: f64,
    /// Sample skewness of the residuals (biased/population form).
    pub skewness: f64,
    /// Sample kurtosis of the residuals (biased/population form; normal ≈ 3).
    pub kurtosis: f64,
}

/// Jarque-Bera test for residual normality, from sample skewness `S` and
/// kurtosis `K`:
///
/// `JB = (n/6)·(S² + (K − 3)²/4)`
///
/// Under normality `JB` is asymptotically χ²(2); the p-value is its upper tail.
/// A large statistic / small p-value is evidence *against* normal residuals.
/// Skewness and kurtosis use the biased (population) moment definitions, matching
/// `statsmodels`.
pub fn jarque_bera(fit: &OlsFit) -> JarqueBera {
    let e = fit.residuals();
    let n = e.len() as f64;
    let mean = e.sum() / n;

    let m2: f64 = e.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / n;
    let m3: f64 = e.iter().map(|v| (v - mean).powi(3)).sum::<f64>() / n;
    let m4: f64 = e.iter().map(|v| (v - mean).powi(4)).sum::<f64>() / n;

    if m2 <= 0.0 {
        return JarqueBera {
            statistic: f64::NAN,
            p_value: f64::NAN,
            skewness: f64::NAN,
            kurtosis: f64::NAN,
        };
    }

    let skewness = m3 / m2.powf(1.5);
    let kurtosis = m4 / (m2 * m2);
    let statistic = (n / 6.0) * (skewness * skewness + (kurtosis - 3.0).powi(2) / 4.0);

    let p_value = match ChiSquared::new(2.0) {
        Ok(dist) if statistic.is_finite() && statistic >= 0.0 => 1.0 - dist.cdf(statistic),
        _ => f64::NAN,
    };

    JarqueBera {
        statistic,
        p_value,
        skewness,
        kurtosis,
    }
}