regression-diagnostics 0.2.0

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

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

/// A generalized linear model fit by **iteratively reweighted least squares**
/// (Fisher scoring) for an exponential-dispersion [`Family`].
///
/// One solver serves every family. Each IRLS step forms the working response
/// `z = η + (y − μ)/(dμ/dη)` and working weights `w = (dμ/dη)² / V(μ)`, then
/// solves the weighted normal equations `(XᵀWX) β = XᵀWz`. At convergence the
/// score equations `Xᵀ(y − μ) · (dμ/dη)/V = 0` hold and `φ·(XᵀWX)⁻¹` is the
/// coefficient covariance, with `φ = 1` for the fixed-dispersion families
/// (Poisson, negative binomial) and the Pearson estimate `φ̂ = χ²/(n − p)` for
/// Gamma.
///
/// See the [module docs](crate::glm) for the response conventions and the link
/// choice.
#[derive(Debug, Clone)]
pub struct GlmFit<F: Family> {
    family: F,
    x: Array2<f64>,
    y: Array1<f64>,
    coefficients: Array1<f64>,
    /// Linear predictor `η = Xβ`.
    eta: Array1<f64>,
    /// Fitted means `μ = g⁻¹(η)`.
    mu: Array1<f64>,
    /// IRLS working weights `wᵢ = (dμ/dη)² / V(μ)` at the MLE.
    weights: Array1<f64>,
    /// Unscaled inverse Fisher information `(XᵀWX)⁻¹` (no dispersion factor).
    cov_unscaled: Array2<f64>,
    /// Dispersion `φ` (`1` when known, else the Pearson estimate).
    dispersion: f64,
    log_likelihood: f64,
    intercept_col: Option<usize>,
    iterations: usize,
    n: usize,
    p: usize,
}

impl<F: Family> GlmFit<F> {
    /// Fit `y ~ X` under `family` (default: up to 100 IRLS iterations, tolerance
    /// `1e-10` on the maximum coefficient step).
    ///
    /// The caller owns the design matrix, including any intercept column, exactly
    /// as with [`OlsFit`](crate::OlsFit) and [`LogisticFit`](crate::logistic::LogisticFit).
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`].
    /// * [`RegressionError::InvalidResponse`] if `y` is outside the family's
    ///   support (negative counts, non-positive Gamma responses).
    /// * [`RegressionError::RankDeficient`] if the weighted design is singular.
    /// * [`RegressionError::NotConverged`] if IRLS fails to converge.
    pub fn new(family: F, x: Array2<f64>, y: Array1<f64>) -> Result<Self> {
        Self::with_options(family, x, y, 100, 1e-10)
    }

    /// Like [`GlmFit::new`] with an explicit iteration cap and step tolerance.
    pub fn with_options(
        family: F,
        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(),
            });
        }
        let y_std = y.as_standard_layout();
        family.validate(y_std.as_slice().expect("standard layout is contiguous"))?;

        let intercept_col = detect_constant_column(&x);

        // Initialize the linear predictor from a smoothed mean, the standard GLM
        // warm start (η⁰ = g(μ⁰)); β is recovered on the first weighted solve.
        let mut eta = Array1::from_shape_fn(n, |i| family.link(family.init_mu(y[i])));
        let mut beta = Array1::<f64>::zeros(p);
        let mut mu = Array1::<f64>::zeros(n);
        let mut weights = Array1::<f64>::zeros(n);
        let mut cov_unscaled = Array2::<f64>::zeros((p, p));
        let mut iterations = 0usize;
        let mut converged = false;

        while iterations < max_iter {
            iterations += 1;

            // Working response z and weights w from the current η.
            let mut z = Array1::<f64>::zeros(n);
            for i in 0..n {
                let mui = family.inverse_link(eta[i]);
                let dmu = family.dmu_deta(eta[i]);
                let v = family.variance(mui);
                mu[i] = mui;
                weights[i] = if v > 0.0 { dmu * dmu / v } else { 0.0 };
                z[i] = eta[i] + (y[i] - mui) / dmu;
            }

            // Weighted normal equations: (XᵀWX) β = XᵀWz.
            let mut xtwx = Array2::<f64>::zeros((p, p));
            let mut xtwz = Array1::<f64>::zeros(p);
            for a in 0..p {
                let mut sz = 0.0;
                for i in 0..n {
                    sz += x[(i, a)] * weights[i] * z[i];
                }
                xtwz[a] = sz;
                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)]);

            let new_beta = inv_arr.dot(&xtwz);
            let step = (&new_beta - &beta)
                .iter()
                .fold(0.0_f64, |m, v| m.max(v.abs()));
            beta = new_beta;
            cov_unscaled = inv_arr;
            eta = x.dot(&beta);

            if !beta.iter().all(|v| v.is_finite()) || beta.iter().any(|v| v.abs() > 1e8) {
                return Err(RegressionError::NotConverged {
                    iterations,
                    msg: "coefficients diverging (perfect fit or unstable link)".into(),
                });
            }
            if step < tol {
                converged = true;
                break;
            }
        }

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

        // Final means, weights, dispersion, and log-likelihood at the MLE.
        for i in 0..n {
            let mui = family.inverse_link(eta[i]);
            let dmu = family.dmu_deta(eta[i]);
            let v = family.variance(mui);
            mu[i] = mui;
            weights[i] = if v > 0.0 { dmu * dmu / v } else { 0.0 };
        }

        let dispersion = if family.dispersion_known() {
            1.0
        } else if n > p {
            // Pearson estimate φ̂ = Σ (yᵢ − μᵢ)²/V(μᵢ) / (n − p).
            let pearson_chi2: f64 = (0..n)
                .map(|i| {
                    let v = family.variance(mu[i]);
                    if v > 0.0 {
                        let r = y[i] - mu[i];
                        r * r / v
                    } else {
                        0.0
                    }
                })
                .sum();
            pearson_chi2 / (n - p) as f64
        } else {
            f64::NAN
        };

        let log_likelihood = (0..n)
            .map(|i| family.loglik(y[i], mu[i], dispersion))
            .sum();

        Ok(Self {
            family,
            x,
            y,
            coefficients: beta,
            eta,
            mu,
            weights,
            cov_unscaled,
            dispersion,
            log_likelihood,
            intercept_col,
            iterations,
            n,
            p,
        })
    }

    /// The family this model was fit under.
    pub fn family(&self) -> &F {
        &self.family
    }

    /// 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 a constant (intercept) column was detected in the design.
    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 response.
    pub fn response(&self) -> ArrayView1<'_, f64> {
        self.y.view()
    }

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

    /// Linear predictor `η = Xβ`.
    pub fn linear_predictor(&self) -> ArrayView1<'_, f64> {
        self.eta.view()
    }

    /// Fitted means `μᵢ = g⁻¹(ηᵢ)`.
    pub fn fitted_means(&self) -> ArrayView1<'_, f64> {
        self.mu.view()
    }

    /// IRLS working weights `wᵢ = (dμ/dη)² / V(μ)` at the MLE.
    pub fn weights(&self) -> ArrayView1<'_, f64> {
        self.weights.view()
    }

    /// Estimated dispersion `φ`: `1` for Poisson/negative binomial, the Pearson
    /// estimate `χ²/(n − p)` for Gamma.
    pub fn dispersion(&self) -> f64 {
        self.dispersion
    }

    /// Maximized log-likelihood (evaluated at the estimated dispersion for
    /// Gamma).
    pub fn log_likelihood(&self) -> f64 {
        self.log_likelihood
    }

    /// Coefficient covariance `φ · (XᵀWX)⁻¹`.
    pub fn covariance(&self) -> Array2<f64> {
        &self.cov_unscaled * self.dispersion
    }

    /// The unscaled inverse Fisher information `(XᵀWX)⁻¹` (no dispersion factor);
    /// this is what the leverage / hat-matrix computation uses.
    pub(crate) fn cov_unscaled(&self) -> ArrayView2<'_, f64> {
        self.cov_unscaled.view()
    }

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

    /// Wald statistics `βⱼ / seⱼ`.
    ///
    /// The reference distribution is the standard normal when the dispersion is
    /// known (Poisson, negative binomial) and Student's *t* with `n − p` degrees
    /// of freedom when it is estimated (Gamma) — see [`GlmFit::p_values`].
    pub fn wald_statistics(&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.
    ///
    /// Uses the standard normal for fixed-dispersion families and Student's *t*
    /// with `n − p` degrees of freedom when the dispersion is estimated (the R
    /// `glm()` convention).
    pub fn p_values(&self) -> Array1<f64> {
        let stat = self.wald_statistics();
        if self.family.dispersion_known() {
            let normal = Normal::new(0.0, 1.0).expect("standard normal");
            Array1::from_shape_fn(self.p, |j| {
                if stat[j].is_finite() {
                    2.0 * (1.0 - normal.cdf(stat[j].abs()))
                } else {
                    f64::NAN
                }
            })
        } else {
            let df = (self.n.saturating_sub(self.p)) as f64;
            let t = StudentsT::new(0.0, 1.0, df.max(1.0)).expect("t distribution");
            Array1::from_shape_fn(self.p, |j| {
                if stat[j].is_finite() {
                    2.0 * (1.0 - t.cdf(stat[j].abs()))
                } else {
                    f64::NAN
                }
            })
        }
    }

    /// Predicted means for a new design matrix `x` (same column layout as the
    /// training design): `μ = g⁻¹(xβ)`.
    pub fn predict_mean(&self, x: ArrayView2<'_, f64>) -> Array1<f64> {
        x.dot(&self.coefficients)
            .mapv(|eta| self.family.inverse_link(eta))
    }
}

/// Detect the first constant column (treated as the intercept). Identical
/// convention to the logistic and OLS fits.
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
}