regression-diagnostics 0.2.0

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

use super::LogisticFit;

/// Result of the Hosmer–Lemeshow goodness-of-fit test.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct HosmerLemeshow {
    /// The Ĥ statistic.
    pub statistic: f64,
    /// Degrees of freedom `g − 2` (`g` = number of groups).
    pub df: usize,
    /// Upper-tail p-value under χ²(df). A **small** p-value indicates **poor**
    /// fit (observed and expected event counts disagree across risk groups).
    pub p_value: f64,
    /// Number of groups actually used.
    pub groups: usize,
}

/// Bundle of overall logistic goodness-of-fit statistics.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct GoodnessOfFit {
    /// Deviance of the intercept-only (null) model, `−2ℓ₀`.
    pub null_deviance: f64,
    /// Residual deviance of the fitted model, `−2ℓ`.
    pub residual_deviance: f64,
    /// Degrees of freedom of the null deviance (`n − 1`).
    pub df_null: f64,
    /// Degrees of freedom of the residual deviance (`n − p`).
    pub df_residual: f64,
    /// McFadden's pseudo-R², `1 − ℓ/ℓ₀`.
    pub mcfadden_r2: f64,
    /// Akaike information criterion, `deviance + 2p`.
    pub aic: f64,
    /// Bayesian information criterion, `deviance + ln(n)·p`.
    pub bic: f64,
}

impl LogisticFit {
    /// Overall goodness-of-fit summary: null/residual deviance, McFadden's
    /// pseudo-R², and AIC/BIC.
    ///
    /// The **residual deviance** is `−2ℓ` (the saturated-model log-likelihood is
    /// zero for ungrouped binary data), and equals the sum of squared deviance
    /// residuals. **McFadden's pseudo-R²** compares the fitted log-likelihood to
    /// the intercept-only model; it is bounded in `[0, 1)` but runs lower than an
    /// OLS R² for comparable fits, so judge it on that scale.
    pub fn goodness_of_fit(&self) -> GoodnessOfFit {
        let n = self.n_observations() as f64;
        let p = self.n_parameters() as f64;
        let y = self.response();

        let residual_deviance = -2.0 * self.log_likelihood();

        // Null model: constant probability = mean(y).
        let ybar = y.sum() / n;
        let ll_null: f64 = y
            .iter()
            .map(|&yi| yi * ybar.ln() + (1.0 - yi) * (1.0 - ybar).ln())
            .sum();
        let null_deviance = -2.0 * ll_null;

        let mcfadden_r2 = if ll_null != 0.0 {
            1.0 - self.log_likelihood() / ll_null
        } else {
            f64::NAN
        };

        GoodnessOfFit {
            null_deviance,
            residual_deviance,
            df_null: n - 1.0,
            df_residual: n - p,
            mcfadden_r2,
            aic: residual_deviance + 2.0 * p,
            bic: residual_deviance + n.ln() * p,
        }
    }

    /// Hosmer–Lemeshow goodness-of-fit test with `groups` risk deciles
    /// (`groups = 10` is the customary choice).
    ///
    /// Observations are ordered by fitted probability and split into `groups`
    /// near-equal bins; the statistic compares observed and expected event counts
    /// per bin,
    ///
    /// `Ĥ = Σ_g (O_g − E_g)² / (n_g · π̄_g · (1 − π̄_g))`,
    ///
    /// which is asymptotically χ²(groups − 2). Unlike most tests here, a **small
    /// p-value means the model fits poorly**. Bins whose mean probability is
    /// exactly 0 or 1 contribute nothing (their variance term is undefined).
    ///
    /// The test is only meaningful when `groups ≥ 3` and there are enough
    /// observations to populate the bins; with fewer than `groups + 1`
    /// observations the statistic is returned as `NaN`.
    pub fn hosmer_lemeshow(&self, groups: usize) -> HosmerLemeshow {
        let n = self.n_observations();
        let y = self.response();
        let p = self.fitted_probabilities();

        if groups < 3 || n < groups + 1 {
            return HosmerLemeshow {
                statistic: f64::NAN,
                df: groups.saturating_sub(2),
                p_value: f64::NAN,
                groups,
            };
        }

        // Order indices by fitted probability.
        let mut order: Vec<usize> = (0..n).collect();
        order.sort_by(|&a, &b| p[a].partial_cmp(&p[b]).unwrap_or(std::cmp::Ordering::Equal));

        let mut statistic = 0.0;
        let mut used_groups = 0usize;
        for g in 0..groups {
            let start = g * n / groups;
            let end = (g + 1) * n / groups;
            if end <= start {
                continue;
            }
            let ng = (end - start) as f64;
            let mut observed = 0.0;
            let mut expected = 0.0;
            for &idx in &order[start..end] {
                observed += y[idx];
                expected += p[idx];
            }
            let pbar = expected / ng;
            if pbar <= 0.0 || pbar >= 1.0 {
                continue;
            }
            let diff = observed - expected;
            statistic += diff * diff / (ng * pbar * (1.0 - pbar));
            used_groups += 1;
        }

        let df = used_groups.saturating_sub(2);
        let p_value = match ChiSquared::new(df as f64) {
            Ok(dist) if df >= 1 && statistic.is_finite() => 1.0 - dist.cdf(statistic),
            _ => f64::NAN,
        };

        HosmerLemeshow {
            statistic,
            df,
            p_value,
            groups: used_groups,
        }
    }
}