regression-diagnostics 0.2.0

use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
use statrs::distribution::{ContinuousCDF, Normal};

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

/// Probability floor/ceiling used to keep IRLS weights away from exactly zero.
const PROB_EPS: f64 = 1e-12;

/// A fitted binary logistic-regression model.
///
/// Fit by **iteratively reweighted least squares** (equivalently, Fisher
/// scoring / Newton–Raphson on the log-likelihood). At convergence the score
/// equations `Xᵀ(y − p) = 0` hold, and the inverse Fisher information
/// `(XᵀWX)⁻¹` — with `W = diag(pᵢ(1−pᵢ))` — gives the coefficient covariance the
/// Wald statistics are built from.
///
/// See the [module docs](crate::logistic) for the response convention and the
/// separation caveat.
#[derive(Debug, Clone)]
pub struct LogisticFit {
    x: Array2<f64>,
    y: Array1<f64>,
    coefficients: Array1<f64>,
    /// Fitted probabilities `pᵢ`.
    probabilities: Array1<f64>,
    /// IRLS weights `wᵢ = pᵢ(1 − pᵢ)`.
    weights: Array1<f64>,
    /// Coefficient covariance `(XᵀWX)⁻¹` at the MLE.
    cov: Array2<f64>,
    log_likelihood: f64,
    intercept_col: Option<usize>,
    iterations: usize,
    n: usize,
    p: usize,
}

impl LogisticFit {
    /// Fit logistic regression of binary `y` on `X` (default: up to 100 IRLS
    /// iterations, tolerance `1e-10` on the coefficient step).
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`].
    /// * [`RegressionError::InvalidResponse`] if `y` is not `0/1` or is all one
    ///   class.
    /// * [`RegressionError::RankDeficient`] if the (weighted) design is singular.
    /// * [`RegressionError::NotConverged`] if IRLS fails to converge (typically
    ///   separation).
    pub fn new(x: Array2<f64>, y: Array1<f64>) -> Result<Self> {
        Self::with_options(x, y, 100, 1e-10)
    }

    /// Like [`LogisticFit::new`] with an explicit iteration cap and tolerance.
    pub fn with_options(x: Array2<f64>, y: Array1<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(),
            });
        }
        // Validate binary response with both classes present.
        let mut saw0 = false;
        let mut saw1 = 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; the fit is not identifiable".into(),
            });
        }

        let intercept_col = detect_constant_column(&x);

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

        while iterations < max_iter {
            iterations += 1;

            // η = Xβ, p = σ(η), w = p(1−p), clamped away from 0/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);
            }

            // Gradient Xᵀ(y − p) and Fisher information XᵀWX.
            let resid = &y - &probabilities;
            let grad = x.t().dot(&resid); // length p
            let mut xtwx = Array2::<f64>::zeros((p, p));
            for a in 0..p {
                for b in a..p {
                    let mut s = 0.0;
                    for i in 0..n {
                        s += x[(i, a)] * weights[i] * x[(i, b)];
                    }
                    xtwx[(a, b)] = s;
                    xtwx[(b, a)] = s;
                }
            }

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

            // Newton step Δ = (XᵀWX)⁻¹ Xᵀ(y − p).
            let delta = inv_arr.dot(&grad);
            beta = &beta + &delta;
            cov = inv_arr;

            let step = delta.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
            if !beta.iter().all(|v| v.is_finite()) || beta.iter().any(|v| v.abs() > 1e8) {
                return Err(RegressionError::NotConverged {
                    iterations,
                    msg: "coefficients diverging (likely perfect separation)".into(),
                });
            }
            if step < tol {
                converged = true;
                break;
            }
        }

        if !converged {
            return Err(RegressionError::NotConverged {
                iterations,
                msg: "IRLS did not reach tolerance (possible quasi-separation)".into(),
            });
        }

        // Final probabilities and log-likelihood at the converged β.
        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 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,
            coefficients: beta,
            probabilities,
            weights,
            cov,
            log_likelihood,
            intercept_col,
            iterations,
            n,
            p,
        })
    }

    /// 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 intercept (constant) column 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()
    }

    /// Estimated 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()
    }

    /// IRLS weights `wᵢ = pᵢ(1 − pᵢ)` at the MLE.
    pub fn weights(&self) -> ArrayView1<'_, f64> {
        self.weights.view()
    }

    /// Coefficient covariance matrix `(XᵀWX)⁻¹`.
    pub fn covariance(&self) -> ArrayView2<'_, f64> {
        self.cov.view()
    }

    /// Maximized log-likelihood.
    pub fn log_likelihood(&self) -> f64 {
        self.log_likelihood
    }

    /// Coefficient standard errors `√diag((XᵀWX)⁻¹)`.
    pub fn coefficient_standard_errors(&self) -> Array1<f64> {
        Array1::from_shape_fn(self.p, |j| self.cov[(j, j)].max(0.0).sqrt())
    }

    /// Wald `z`-statistics `βⱼ / seⱼ`.
    pub fn z_values(&self) -> Array1<f64> {
        let se = self.coefficient_standard_errors();
        Array1::from_shape_fn(self.p, |j| {
            if se[j] > 0.0 {
                self.coefficients[j] / se[j]
            } else {
                f64::NAN
            }
        })
    }

    /// Two-sided Wald p-values from the standard normal.
    pub fn p_values(&self) -> Array1<f64> {
        let z = self.z_values();
        let normal = Normal::new(0.0, 1.0).expect("standard normal");
        Array1::from_shape_fn(self.p, |j| {
            if z[j].is_finite() {
                2.0 * (1.0 - normal.cdf(z[j].abs()))
            } else {
                f64::NAN
            }
        })
    }

    /// Predicted probabilities for a new design matrix `x` (same column layout
    /// as the training design).
    pub fn predict_proba(&self, x: ArrayView2<'_, f64>) -> Array1<f64> {
        x.dot(&self.coefficients).mapv(sigmoid)
    }
}

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

/// Detect the first constant column (treated as the intercept).
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
}