regression-diagnostics 0.2.0

use ndarray::{Array1, Array2, ArrayView1, ArrayView2};

use crate::error::{RegressionError, Result};
use crate::linalg::dmatrix_from_rows;

const PROB_EPS: f64 = 1e-12;

/// A fitted **ridge-penalized logistic regression** — a penalized GLM — and its
/// shrinkage-aware diagnostics.
///
/// Maximizes the L2-penalized log-likelihood
///
/// `ℓ(β) − ½λ‖β_pen‖²`
///
/// (a detected intercept column is left unpenalized), fit by **penalized IRLS**:
/// each Newton step solves `(XᵀWX + λP) Δ = Xᵀ(y − p) − λPβ` with
/// `W = diag(pᵢ(1 − pᵢ))` and `P` the penalty selector. The penalty is the GLM
/// analogue of ridge — it tames separation and multicollinearity in logistic
/// regression, at the cost of biased-but-lower-variance coefficients.
///
/// # The diagnostic that changes: effective degrees of freedom
///
/// As in ridge OLS, the penalty means the model no longer spends `p` degrees of
/// freedom. The **effective df** is the trace of the penalized hat matrix,
///
/// `df = tr[ (XᵀWX + λP)⁻¹ XᵀWX ]`,
///
/// which falls from `p` toward the unpenalized/intercept count as `λ` grows. It
/// replaces the raw parameter count in [`aic`](Self::aic) / [`bic`](Self::bic),
/// and the coefficient covariance is the **sandwich**
/// `(XᵀWX + λP)⁻¹ (XᵀWX) (XᵀWX + λP)⁻¹`, not the naïve inverse information — so
/// the reported standard errors account for the shrinkage.
#[derive(Debug, Clone)]
pub struct PenalizedLogisticFit {
    x: Array2<f64>,
    y: Array1<f64>,
    lambda: f64,
    coefficients: Array1<f64>,
    probabilities: Array1<f64>,
    cov: Array2<f64>,
    log_likelihood: f64,
    effective_df: f64,
    intercept_col: Option<usize>,
    iterations: usize,
    n: usize,
    p: usize,
}

impl PenalizedLogisticFit {
    /// Fit ridge-penalized logistic regression of binary `y` on `X` with penalty
    /// `lambda ≥ 0` (default: up to 100 IRLS iterations, tolerance `1e-10`).
    ///
    /// At `lambda = 0` this reproduces the ordinary
    /// [`LogisticFit`](crate::logistic::LogisticFit).
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`].
    /// * [`RegressionError::InvalidParameter`] if `lambda < 0`.
    /// * [`RegressionError::InvalidResponse`] if `y` is not `0/1` or is one class.
    /// * [`RegressionError::RankDeficient`] if the penalized information is
    ///   singular.
    /// * [`RegressionError::NotConverged`] if IRLS fails to converge.
    pub fn new(x: Array2<f64>, y: Array1<f64>, lambda: f64) -> Result<Self> {
        Self::with_options(x, y, lambda, 100, 1e-10)
    }

    /// Like [`PenalizedLogisticFit::new`] with an explicit iteration cap and
    /// tolerance.
    pub fn with_options(
        x: Array2<f64>,
        y: Array1<f64>,
        lambda: f64,
        max_iter: usize,
        tol: f64,
    ) -> Result<Self> {
        let n = x.nrows();
        let p = x.ncols();
        if n == 0 || p == 0 {
            return Err(RegressionError::EmptyInput { what: "X" });
        }
        if y.len() != n {
            return Err(RegressionError::ShapeMismatch {
                what: "y length vs X rows",
                expected: n,
                got: y.len(),
            });
        }
        if lambda < 0.0 || lambda.is_nan() {
            return Err(RegressionError::InvalidParameter {
                msg: format!("penalty lambda must be >= 0, got {lambda}"),
            });
        }
        let (mut saw0, mut saw1) = (false, false);
        for &v in y.iter() {
            if v == 0.0 {
                saw0 = true;
            } else if v == 1.0 {
                saw1 = true;
            } else {
                return Err(RegressionError::InvalidResponse {
                    msg: format!("response must be 0 or 1, found {v}"),
                });
            }
        }
        if !(saw0 && saw1) {
            return Err(RegressionError::InvalidResponse {
                msg: "response is entirely one class".into(),
            });
        }

        let intercept_col = detect_constant_column(&x);
        // Penalty selector: 1 for penalized columns, 0 for the intercept.
        let pen: Vec<f64> = (0..p)
            .map(|j| if Some(j) == intercept_col { 0.0 } else { 1.0 })
            .collect();

        let mut beta = Array1::<f64>::zeros(p);
        let mut probabilities = Array1::<f64>::zeros(n);
        let mut weights = Array1::<f64>::zeros(n);
        let mut xtwx_pen_inv = Array2::<f64>::zeros((p, p));
        let mut iterations = 0usize;
        let mut converged = false;

        while iterations < max_iter {
            iterations += 1;
            let eta = x.dot(&beta);
            for i in 0..n {
                let pi = sigmoid(eta[i]).clamp(PROB_EPS, 1.0 - PROB_EPS);
                probabilities[i] = pi;
                weights[i] = pi * (1.0 - pi);
            }

            // Penalized gradient g = Xᵀ(y − p) − λ P β.
            let resid = &y - &probabilities;
            let mut grad = x.t().dot(&resid);
            for j in 0..p {
                grad[j] -= lambda * pen[j] * beta[j];
            }
            // Penalized information A = XᵀWX + λ P.
            let mut a = Array2::<f64>::zeros((p, p));
            for r in 0..p {
                for c in r..p {
                    let mut s = 0.0;
                    for i in 0..n {
                        s += x[(i, r)] * weights[i] * x[(i, c)];
                    }
                    a[(r, c)] = s;
                    a[(c, r)] = s;
                }
            }
            for j in 0..p {
                a[(j, j)] += lambda * pen[j];
            }

            let a_dm = dmatrix_from_rows(p, p, a.as_standard_layout().as_slice().unwrap());
            let inv = a_dm.try_inverse().ok_or(RegressionError::RankDeficient)?;
            let inv_arr = Array2::from_shape_fn((p, p), |(i, j)| inv[(i, j)]);

            let delta = inv_arr.dot(&grad);
            beta = &beta + &delta;
            xtwx_pen_inv = inv_arr;

            let step = delta.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
            if !beta.iter().all(|v| v.is_finite()) {
                return Err(RegressionError::NotConverged {
                    iterations,
                    msg: "coefficients diverging".into(),
                });
            }
            if step < tol {
                converged = true;
                break;
            }
        }
        if !converged {
            return Err(RegressionError::NotConverged {
                iterations,
                msg: "penalized IRLS did not reach tolerance".into(),
            });
        }

        // Final probabilities, weights, and the unpenalized XᵀWX.
        let eta = x.dot(&beta);
        for i in 0..n {
            let pi = sigmoid(eta[i]).clamp(PROB_EPS, 1.0 - PROB_EPS);
            probabilities[i] = pi;
            weights[i] = pi * (1.0 - pi);
        }
        let mut xtwx = Array2::<f64>::zeros((p, p));
        for r in 0..p {
            for c in r..p {
                let mut s = 0.0;
                for i in 0..n {
                    s += x[(i, r)] * weights[i] * x[(i, c)];
                }
                xtwx[(r, c)] = s;
                xtwx[(c, r)] = s;
            }
        }
        // effective df = tr[(XᵀWX + λP)⁻¹ XᵀWX].
        let effective_df = (0..p)
            .map(|i| (0..p).map(|kk| xtwx_pen_inv[(i, kk)] * xtwx[(kk, i)]).sum::<f64>())
            .sum();
        // Sandwich covariance A⁻¹ (XᵀWX) A⁻¹.
        let mid = xtwx.dot(&xtwx_pen_inv);
        let cov = xtwx_pen_inv.dot(&mid);

        let log_likelihood = (0..n)
            .map(|i| {
                let pi = probabilities[i];
                y[i] * pi.ln() + (1.0 - y[i]) * (1.0 - pi).ln()
            })
            .sum();

        Ok(Self {
            x,
            y,
            lambda,
            coefficients: beta,
            probabilities,
            cov,
            log_likelihood,
            effective_df,
            intercept_col,
            iterations,
            n,
            p,
        })
    }

    /// The penalty `λ` this model was fit with.
    pub fn lambda(&self) -> f64 {
        self.lambda
    }

    /// Number of observations.
    pub fn n_observations(&self) -> usize {
        self.n
    }

    /// Number of coefficients (design columns, intercept included).
    pub fn n_parameters(&self) -> usize {
        self.p
    }

    /// Whether an unpenalized intercept is present.
    pub fn has_intercept(&self) -> bool {
        self.intercept_col.is_some()
    }

    /// IRLS iterations taken to converge.
    pub fn iterations(&self) -> usize {
        self.iterations
    }

    /// The design matrix as fitted.
    pub fn design_matrix(&self) -> ArrayView2<'_, f64> {
        self.x.view()
    }

    /// The binary response.
    pub fn response(&self) -> ArrayView1<'_, f64> {
        self.y.view()
    }

    /// Penalized coefficients (log-odds scale), aligned to the design columns.
    pub fn coefficients(&self) -> ArrayView1<'_, f64> {
        self.coefficients.view()
    }

    /// Fitted probabilities `pᵢ = P(yᵢ = 1)`.
    pub fn fitted_probabilities(&self) -> ArrayView1<'_, f64> {
        self.probabilities.view()
    }

    /// Sandwich coefficient covariance `(XᵀWX + λP)⁻¹ (XᵀWX) (XᵀWX + λP)⁻¹`.
    pub fn covariance(&self) -> ArrayView2<'_, f64> {
        self.cov.view()
    }

    /// Unpenalized log-likelihood at the penalized estimate.
    pub fn log_likelihood(&self) -> f64 {
        self.log_likelihood
    }

    /// Effective degrees of freedom `tr[(XᵀWX + λP)⁻¹ XᵀWX]` — falls from `p`
    /// toward the intercept count as `λ` grows.
    pub fn effective_df(&self) -> f64 {
        self.effective_df
    }

    /// Coefficient standard errors from the sandwich covariance.
    pub fn coefficient_standard_errors(&self) -> Array1<f64> {
        Array1::from_shape_fn(self.p, |j| self.cov[(j, j)].max(0.0).sqrt())
    }

    /// Residual deviance `−2ℓ` (unpenalized log-likelihood).
    pub fn residual_deviance(&self) -> f64 {
        -2.0 * self.log_likelihood
    }

    /// AIC using the effective degrees of freedom, `−2ℓ + 2·df`.
    pub fn aic(&self) -> f64 {
        -2.0 * self.log_likelihood + 2.0 * self.effective_df
    }

    /// BIC using the effective degrees of freedom, `−2ℓ + ln(n)·df`.
    pub fn bic(&self) -> f64 {
        -2.0 * self.log_likelihood + (self.n as f64).ln() * self.effective_df
    }
}

fn sigmoid(z: f64) -> f64 {
    if z >= 0.0 {
        1.0 / (1.0 + (-z).exp())
    } else {
        let e = z.exp();
        e / (1.0 + e)
    }
}

fn detect_constant_column(x: &Array2<f64>) -> Option<usize> {
    for (j, col) in x.columns().into_iter().enumerate() {
        let first = col[0];
        let scale = first.abs().max(1.0);
        if col.iter().all(|&v| (v - first).abs() <= 1e-12 * scale) {
            return Some(j);
        }
    }
    None
}