regression-diagnostics 0.2.0

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

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

/// A fitted ridge-regression model and its diagnostics.
///
/// Ridge solves `min ‖y − Xβ‖² + λ‖β_pen‖²`. This type fits it in **closed form
/// via the SVD** of the (centered) predictors, which is both numerically stable
/// and hands us the singular values that every ridge diagnostic is expressed in.
///
/// # Intercept and scaling
///
/// If a constant column is detected it is treated as an **unpenalized
/// intercept**: predictors and response are mean-centered, the penalty is
/// applied only to the slopes, and the intercept is recovered from the means
/// (the same convention as scikit-learn's `Ridge`). Ridge is not scale-invariant,
/// so standardize predictors beforehand if that matters for your `λ`.
///
/// # What the diagnostics mean
///
/// The OLS notions still exist but take their ridge forms: leverage is the
/// diagonal of `H_λ`, and the parameter count is the **effective degrees of
/// freedom** `df(λ) = Σ dⱼ²/(dⱼ²+λ)` (plus one for the intercept), which slides
/// smoothly from `p` at `λ = 0` down toward `1` as `λ → ∞`.
#[derive(Debug, Clone)]
pub struct RidgeFit {
    x: Array2<f64>,
    y: Array1<f64>,
    lambda: f64,
    coefficients: Array1<f64>,
    fitted: Array1<f64>,
    residuals: Array1<f64>,
    leverage: Array1<f64>,
    /// Singular values of the centered predictor matrix.
    singular_values: Vec<f64>,
    effective_df: f64,
    rss: f64,
    intercept_col: Option<usize>,
    n: usize,
    p: usize,
}

impl RidgeFit {
    /// Fit ridge regression of `y` on `X` with penalty `lambda ≥ 0`.
    ///
    /// A constant column of `X`, if present, is auto-detected and used as an
    /// unpenalized intercept (see the [type docs](RidgeFit#intercept-and-scaling)).
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`] as
    ///   for [`OlsFit`](crate::OlsFit).
    /// * [`RegressionError::InvalidParameter`] if `lambda < 0`.
    ///
    /// Unlike OLS, ridge is well-defined even when `n ≤ p` or the predictors are
    /// collinear (that is much of the point), so those are not errors here.
    pub fn new(x: Array2<f64>, y: Array1<f64>, lambda: f64) -> Result<Self> {
        if x.nrows() == 0 || x.ncols() == 0 {
            return Err(RegressionError::EmptyInput { what: "X" });
        }
        if y.len() != x.nrows() {
            return Err(RegressionError::ShapeMismatch {
                what: "y length vs X rows",
                expected: x.nrows(),
                got: y.len(),
            });
        }
        if lambda < 0.0 || lambda.is_nan() {
            return Err(RegressionError::InvalidParameter {
                msg: format!("ridge lambda must be >= 0, got {lambda}"),
            });
        }

        let n = x.nrows();
        let p = x.ncols();
        let intercept_col = detect_constant_column(&x);

        // Predictor (non-intercept) column indices.
        let pred: Vec<usize> = (0..p).filter(|&j| Some(j) != intercept_col).collect();
        let q = pred.len();

        // Center predictors and response when there is an intercept; otherwise
        // penalize everything on its raw scale.
        let has_intercept = intercept_col.is_some();
        let y_mean = if has_intercept {
            y.sum() / n as f64
        } else {
            0.0
        };
        let mut x_means = vec![0.0; q];
        if has_intercept {
            for (k, &j) in pred.iter().enumerate() {
                x_means[k] = x.column(j).sum() / n as f64;
            }
        }

        // Build the centered predictor matrix (n × q) in row-major order.
        let mut xc = vec![0.0; n * q];
        for i in 0..n {
            for (k, &j) in pred.iter().enumerate() {
                xc[i * q + k] = x[(i, j)] - x_means[k];
            }
        }
        let yc: Vec<f64> = (0..n).map(|i| y[i] - y_mean).collect();

        // Thin SVD of the centered predictors: Xc = U S Vᵀ.
        let xc_dm = dmatrix_from_rows(n, q, &xc);
        let svd = xc_dm.svd(true, true);
        let u = svd.u.ok_or(RegressionError::RankDeficient)?; // n × q
        let s = svd.singular_values; // length min(n, q)
        let v_t = svd.v_t.ok_or(RegressionError::RankDeficient)?; // (min) × q
        let r = s.len();

        // a = Uᵀ yc
        let a: Vec<f64> = (0..r)
            .map(|j| (0..n).map(|i| u[(i, j)] * yc[i]).sum::<f64>())
            .collect();
        // filter factors f_j = d_j / (d_j² + λ)
        let filt: Vec<f64> = (0..r).map(|j| s[j] / (s[j] * s[j] + lambda)).collect();
        // slopes β_p[k] = Σ_j V[k,j] f_j a_j = Σ_j v_t[j,k] f_j a_j
        let mut slopes = vec![0.0; q];
        for (k, slope) in slopes.iter_mut().enumerate() {
            *slope = (0..r).map(|j| v_t[(j, k)] * filt[j] * a[j]).sum();
        }

        // Assemble coefficients aligned to the original columns.
        let mut coefficients = Array1::<f64>::zeros(p);
        for (k, &j) in pred.iter().enumerate() {
            coefficients[j] = slopes[k];
        }
        if let Some(c) = intercept_col {
            let intercept = y_mean - (0..q).map(|k| x_means[k] * slopes[k]).sum::<f64>();
            coefficients[c] = intercept;
        }

        // Fitted / residuals from the original design.
        let fitted = x.dot(&coefficients);
        let residuals = &y - &fitted;
        let rss: f64 = residuals.iter().map(|e| e * e).sum();

        // Ridge leverage (diagonal of H_λ) and effective df.
        // shrink_j = d_j² / (d_j² + λ)
        let shrink: Vec<f64> = (0..r)
            .map(|j| {
                let d2 = s[j] * s[j];
                d2 / (d2 + lambda)
            })
            .collect();
        let base = if has_intercept { 1.0 / n as f64 } else { 0.0 };
        let base_df = if has_intercept { 1.0 } else { 0.0 };
        let leverage = Array1::from_shape_fn(n, |i| {
            base + (0..r)
                .map(|j| u[(i, j)] * u[(i, j)] * shrink[j])
                .sum::<f64>()
        });
        let effective_df = base_df + shrink.iter().sum::<f64>();

        Ok(Self {
            x,
            y,
            lambda,
            coefficients,
            fitted,
            residuals,
            leverage,
            singular_values: s.iter().copied().collect(),
            effective_df,
            rss,
            intercept_col,
            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 intercept (constant column) is present and unpenalized.
    pub fn has_intercept(&self) -> bool {
        self.intercept_col.is_some()
    }

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

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

    /// Singular values of the centered predictor matrix — the `dⱼ` that the
    /// shrinkage factors `dⱼ²/(dⱼ²+λ)` and the effective degrees of freedom are
    /// expressed in.
    pub fn singular_values(&self) -> &[f64] {
        &self.singular_values
    }

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

    /// Fitted values `ŷ = Xβ`.
    pub fn fitted_values(&self) -> ArrayView1<'_, f64> {
        self.fitted.view()
    }

    /// Residuals `y − ŷ`.
    pub fn residuals(&self) -> ArrayView1<'_, f64> {
        self.residuals.view()
    }

    /// Residual sum of squares.
    pub fn residual_sum_of_squares(&self) -> f64 {
        self.rss
    }

    /// Ridge leverage — the diagonal of `H_λ = X(XᵀX + λI)⁻¹Xᵀ` — computed from
    /// the SVD without ever forming the `n × n` hat matrix.
    ///
    /// Unlike OLS leverage these do **not** sum to the number of columns; they
    /// sum to the effective degrees of freedom [`effective_df`](Self::effective_df),
    /// which is the ridge analogue of that identity.
    pub fn leverage(&self) -> ArrayView1<'_, f64> {
        self.leverage.view()
    }

    /// Effective degrees of freedom `df(λ) = base + Σ dⱼ²/(dⱼ²+λ)`, where `base`
    /// is `1` for the unpenalized intercept (else `0`).
    ///
    /// This is the parameter count ridge actually spends: it equals `p` at
    /// `λ = 0` and shrinks toward `1` (just the intercept) as `λ → ∞`. It drives
    /// the effective residual df and the information criteria below.
    pub fn effective_df(&self) -> f64 {
        self.effective_df
    }

    /// Effective residual degrees of freedom `n − df(λ)`.
    pub fn effective_residual_df(&self) -> f64 {
        self.n as f64 - self.effective_df
    }

    /// Effective residual variance estimate `RSS / (n − df(λ))`.
    pub fn residual_variance(&self) -> f64 {
        self.rss / self.effective_residual_df()
    }

    /// Generalized Cross-Validation score
    /// `GCV(λ) = (RSS / n) / (1 − df(λ)/n)²`.
    ///
    /// A rotation-invariant approximation to leave-one-out CV; the `λ` minimizing
    /// it is a standard, data-driven penalty choice (see
    /// [`select_lambda_gcv`]).
    pub fn gcv(&self) -> f64 {
        let denom = 1.0 - self.effective_df / self.n as f64;
        if denom <= 0.0 {
            return f64::INFINITY;
        }
        (self.rss / self.n as f64) / (denom * denom)
    }

    /// Gaussian log-likelihood at the fitted residual variance (same form as the
    /// OLS log-likelihood, using `RSS` from the ridge fit).
    pub fn log_likelihood(&self) -> f64 {
        let n = self.n as f64;
        -0.5 * n * ((2.0 * std::f64::consts::PI).ln() + 1.0 + (self.rss / n).ln())
    }

    /// AIC using the **effective** degrees of freedom as the parameter count:
    /// `AIC = −2ℓ + 2·df(λ)`.
    ///
    /// Using `df(λ)` rather than `p` is what makes information criteria
    /// meaningful under shrinkage — the model is charged for the degrees of
    /// freedom it effectively uses, not the nominal column count.
    pub fn aic(&self) -> f64 {
        -2.0 * self.log_likelihood() + 2.0 * self.effective_df
    }

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

    /// Ridge Variance Inflation Factors, aligned to the design columns
    /// (intercept slot is `NaN`).
    ///
    /// Computed as the diagonal of `(R + λ_c I)⁻¹ R (R + λ_c I)⁻¹` on the
    /// **standardized** predictors, where `R` is their correlation matrix. This
    /// is the exact variance-inflation of the ridge coefficient estimates and it
    /// **reduces to the ordinary OLS VIF at `λ = 0`**, which is what makes the
    /// "VIF before vs after regularization" comparison meaningful: as `λ` grows
    /// these fall, quantifying how ridge tames collinearity. `λ_c` is the penalty
    /// on the correlation scale (`λ` divided by `n`, since the correlation matrix
    /// uses the `1/n`-scaled cross-products).
    pub fn ridge_vif(&self) -> Vec<f64> {
        let pred: Vec<usize> = (0..self.p)
            .filter(|&j| Some(j) != self.intercept_col)
            .collect();
        let q = pred.len();
        let mut out = vec![f64::NAN; self.p];
        if q == 0 {
            return out;
        }

        // Correlation matrix R of the predictors.
        let n = self.n as f64;
        let mut means = vec![0.0; q];
        let mut sds = vec![0.0; q];
        for (k, &j) in pred.iter().enumerate() {
            let col = self.x.column(j);
            let m = col.sum() / n;
            means[k] = m;
            sds[k] = (col.iter().map(|v| (v - m).powi(2)).sum::<f64>() / n).sqrt();
        }
        let corr = |k: usize, l: usize| -> f64 {
            if sds[k] <= 0.0 || sds[l] <= 0.0 {
                return if k == l { 1.0 } else { 0.0 };
            }
            let (jk, jl) = (pred[k], pred[l]);
            let ck = self.x.column(jk);
            let cl = self.x.column(jl);
            let cov: f64 = (0..self.n)
                .map(|i| (ck[i] - means[k]) * (cl[i] - means[l]))
                .sum::<f64>()
                / n;
            cov / (sds[k] * sds[l])
        };

        let lambda_c = self.lambda / n;
        // Build (R + λ_c I).
        let mut a = Array2::<f64>::zeros((q, q));
        let mut rmat = Array2::<f64>::zeros((q, q));
        for k in 0..q {
            for l in 0..q {
                let r = corr(k, l);
                rmat[(k, l)] = r;
                a[(k, l)] = r + if k == l { lambda_c } else { 0.0 };
            }
        }
        // M = A⁻¹ R A⁻¹
        let a_dm = dmatrix_from_rows(q, q, a.as_standard_layout().as_slice().unwrap());
        let a_inv = match a_dm.try_inverse() {
            Some(inv) => inv,
            None => return out,
        };
        let r_dm = dmatrix_from_rows(q, q, rmat.as_standard_layout().as_slice().unwrap());
        let m = &a_inv * r_dm * &a_inv;
        for (k, &j) in pred.iter().enumerate() {
            out[j] = m[(k, k)];
        }
        out
    }
}

/// Select the ridge penalty that minimizes GCV over a grid of candidate `λ`s,
/// returning the best [`RidgeFit`].
///
/// Refits at each candidate (each fit is a single SVD-based closed-form solve)
/// and keeps the one with the smallest [`RidgeFit::gcv`]. The grid is the
/// caller's to choose — a geometric sweep such as `10.^{-3..3}` is typical.
///
/// # Errors
///
/// [`RegressionError::InvalidParameter`] if `lambdas` is empty; otherwise any
/// error from [`RidgeFit::new`].
pub fn select_lambda_gcv(x: Array2<f64>, y: Array1<f64>, lambdas: &[f64]) -> Result<RidgeFit> {
    if lambdas.is_empty() {
        return Err(RegressionError::InvalidParameter {
            msg: "lambda grid must be non-empty".into(),
        });
    }
    let mut best: Option<RidgeFit> = None;
    for &lam in lambdas {
        let fit = RidgeFit::new(x.clone(), y.clone(), lam)?;
        let better = match &best {
            None => true,
            Some(b) => fit.gcv() < b.gcv(),
        };
        if better {
            best = Some(fit);
        }
    }
    Ok(best.expect("non-empty grid yields a fit"))
}

/// 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
}