regression-diagnostics 0.2.0

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

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

/// A fitted **elastic-net** regression model and its diagnostics.
///
/// Elastic net minimizes
///
/// `(1/2n)‖y − Xβ‖² + λ[ α‖β_pen‖₁ + ½(1 − α)‖β_pen‖² ]`,
///
/// blending the lasso (`α = 1`) and ridge (`α → 0`) penalties. The `α` mixing
/// parameter controls sparsity-versus-grouping: pure lasso arbitrarily picks one
/// of a set of correlated predictors, while a little ridge (`α < 1`) shares the
/// coefficient across the group. Fit by **cyclic coordinate descent** with
/// soft-thresholding (glmnet-style).
///
/// # Intercept and scaling
///
/// Identical conventions to [`LassoFit`](super::LassoFit): a detected constant
/// column is an **unpenalized intercept**, predictors are **standardized
/// internally**, and `λ` is on the standardized `(1/2n)`-objective scale.
///
/// # Effective degrees of freedom
///
/// Unlike lasso, where the degrees of freedom are just the active-set size, the
/// ridge part shrinks the surviving coefficients, so the elastic net spends
/// *fewer* than `|active set|` degrees of freedom. This type reports the
/// shrinkage-aware trace
///
/// `df = tr[ Z_A (Z_Aᵀ Z_A + n·λ(1 − α) I)⁻¹ Z_Aᵀ ] (+1 for the intercept)`,
///
/// over the standardized active columns `Z_A`. It reduces to `|active set|` at
/// `α = 1` (recovering the lasso count) and to the full ridge effective df when
/// nothing is zeroed. The information criteria use it.
#[derive(Debug, Clone)]
pub struct ElasticNetFit {
    x: Array2<f64>,
    y: Array1<f64>,
    lambda: f64,
    alpha: f64,
    coefficients: Array1<f64>,
    fitted: Array1<f64>,
    residuals: Array1<f64>,
    rss: f64,
    intercept_col: Option<usize>,
    n_nonzero: usize,
    effective_df: f64,
    iterations: usize,
    n: usize,
    p: usize,
}

impl ElasticNetFit {
    /// Fit elastic-net regression of `y` on `X` with penalty `lambda ≥ 0` and
    /// mixing `alpha ∈ [0, 1]` (default tolerance `1e-7`, up to `10_000` sweeps).
    ///
    /// `alpha = 1` is pure lasso (equivalent to [`LassoFit`](super::LassoFit));
    /// `alpha = 0` is pure ridge on the standardized scale.
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`].
    /// * [`RegressionError::InvalidParameter`] if `lambda < 0` or `alpha ∉ [0, 1]`.
    /// * [`RegressionError::NotConverged`] if coordinate descent does not
    ///   converge within the iteration budget.
    pub fn new(x: Array2<f64>, y: Array1<f64>, lambda: f64, alpha: f64) -> Result<Self> {
        Self::with_options(x, y, lambda, alpha, 1e-7, 10_000)
    }

    /// Like [`ElasticNetFit::new`] with an explicit tolerance and sweep cap.
    pub fn with_options(
        x: Array2<f64>,
        y: Array1<f64>,
        lambda: f64,
        alpha: f64,
        tol: f64,
        max_iter: usize,
    ) -> 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!("elastic-net lambda must be >= 0, got {lambda}"),
            });
        }
        if !(0.0..=1.0).contains(&alpha) {
            return Err(RegressionError::InvalidParameter {
                msg: format!("elastic-net alpha must be in [0, 1], got {alpha}"),
            });
        }

        let n = x.nrows();
        let p = x.ncols();
        let intercept_col = detect_constant_column(&x);
        let has_intercept = intercept_col.is_some();
        let pred: Vec<usize> = (0..p).filter(|&j| Some(j) != intercept_col).collect();
        let q = pred.len();

        let nf = n as f64;
        let y_mean = if has_intercept { y.sum() / nf } else { 0.0 };

        // Standardize predictors (population sd), exactly as lasso does.
        let mut means = vec![0.0; q];
        let mut sds = vec![1.0; q];
        for (k, &j) in pred.iter().enumerate() {
            let col = x.column(j);
            let m = if has_intercept { col.sum() / nf } else { 0.0 };
            means[k] = m;
            let sd = (col.iter().map(|v| (v - m).powi(2)).sum::<f64>() / nf).sqrt();
            sds[k] = if sd > 0.0 { sd } else { 1.0 };
        }
        let mut z = vec![0.0f64; n * q];
        for i in 0..n {
            for (k, &j) in pred.iter().enumerate() {
                z[i * q + k] = (x[(i, j)] - means[k]) / sds[k];
            }
        }
        let yc: Vec<f64> = (0..n).map(|i| y[i] - y_mean).collect();

        // Coordinate descent with the elastic-net update:
        //   β_k = soft(ρ_k, λα) / (1 + λ(1−α)),  ρ_k = (1/n) z_kᵀ r + β_k.
        let l1 = lambda * alpha;
        let l2 = lambda * (1.0 - alpha);
        let mut beta = vec![0.0f64; q];
        let mut r = yc.clone();
        let mut iterations = 0usize;
        let mut converged = false;
        while iterations < max_iter {
            iterations += 1;
            let mut max_delta = 0.0f64;
            for k in 0..q {
                let mut zr = 0.0;
                for i in 0..n {
                    zr += z[i * q + k] * r[i];
                }
                let rho = zr / nf + beta[k];
                let new = soft_threshold(rho, l1) / (1.0 + l2);
                let delta = new - beta[k];
                if delta != 0.0 {
                    for i in 0..n {
                        r[i] -= z[i * q + k] * delta;
                    }
                    beta[k] = new;
                    max_delta = max_delta.max(delta.abs());
                }
            }
            if max_delta < tol {
                converged = true;
                break;
            }
        }
        if !converged {
            return Err(RegressionError::NotConverged {
                iterations,
                msg: "elastic-net coordinate descent did not reach tolerance".into(),
            });
        }

        // Back to original scale.
        let mut coefficients = Array1::<f64>::zeros(p);
        let mut slopes_orig = vec![0.0; q];
        for (k, &j) in pred.iter().enumerate() {
            let b = beta[k] / sds[k];
            slopes_orig[k] = b;
            coefficients[j] = b;
        }
        if let Some(c) = intercept_col {
            coefficients[c] = y_mean - (0..q).map(|k| means[k] * slopes_orig[k]).sum::<f64>();
        }

        let fitted = x.dot(&coefficients);
        let residuals = &y - &fitted;
        let rss: f64 = residuals.iter().map(|e| e * e).sum();

        // Active set (standardized indices).
        let active: Vec<usize> = (0..q).filter(|&k| beta[k].abs() > 1e-12).collect();
        let n_nonzero = active.len();
        let effective_df = elastic_net_df(&z, n, q, &active, nf * l2)
            + if has_intercept { 1.0 } else { 0.0 };

        Ok(Self {
            x,
            y,
            lambda,
            alpha,
            coefficients,
            fitted,
            residuals,
            rss,
            intercept_col,
            n_nonzero,
            effective_df,
            iterations,
            n,
            p,
        })
    }

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

    /// The mixing parameter `α` (1 = lasso, 0 = ridge).
    pub fn alpha(&self) -> f64 {
        self.alpha
    }

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

    /// Coordinate-descent sweeps 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()
    }

    /// Elastic-net coefficients, aligned to the design columns. Coefficients
    /// shrunk out are exactly `0.0`.
    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
    }

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

    /// Indices of the design columns with non-zero coefficients — the **active
    /// set** (excludes the intercept).
    pub fn active_set(&self) -> Vec<usize> {
        (0..self.p)
            .filter(|&j| Some(j) != self.intercept_col && self.coefficients[j].abs() > 1e-12)
            .collect()
    }

    /// Number of non-zero penalized coefficients (active-set size).
    pub fn n_nonzero(&self) -> usize {
        self.n_nonzero
    }

    /// Shrinkage-aware **effective degrees of freedom** (see the type docs):
    /// the active-set trace under the ridge part, plus one for the intercept.
    pub fn effective_df(&self) -> f64 {
        self.effective_df
    }

    /// Gaussian log-likelihood at the fitted residual variance.
    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.
    pub fn aic(&self) -> f64 {
        -2.0 * self.log_likelihood() + 2.0 * self.effective_df
    }

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

/// Trace of the ridge hat matrix on the standardized active columns:
/// `tr[ Z_A (Z_Aᵀ Z_A + ridge·I)⁻¹ Z_Aᵀ ] = tr[ (Z_AᵀZ_A + ridge·I)⁻¹ Z_AᵀZ_A ]`,
/// where `ridge = n·λ(1−α)`. Equals `|A|` when `ridge = 0`.
fn elastic_net_df(z: &[f64], n: usize, q: usize, active: &[usize], ridge: f64) -> f64 {
    let a = active.len();
    if a == 0 {
        return 0.0;
    }
    if ridge <= 0.0 {
        return a as f64; // projection onto the active columns
    }
    // Gram matrix G = Z_AᵀZ_A over the active columns.
    let mut g = Array2::<f64>::zeros((a, a));
    for (ii, &ki) in active.iter().enumerate() {
        for (jj, &kj) in active.iter().enumerate() {
            let mut s = 0.0;
            for row in 0..n {
                s += z[row * q + ki] * z[row * q + kj];
            }
            g[(ii, jj)] = s;
        }
    }
    // M = G + ridge·I.
    let mut m = g.clone();
    for d in 0..a {
        m[(d, d)] += ridge;
    }
    let m_dm = dmatrix_from_rows(a, a, m.as_standard_layout().as_slice().unwrap());
    let inv = match m_dm.try_inverse() {
        Some(inv) => inv,
        None => return a as f64,
    };
    // trace(M⁻¹ G) = Σ_ij inv[i,j] · G[j,i].
    let mut tr = 0.0;
    for i in 0..a {
        for j in 0..a {
            tr += inv[(i, j)] * g[(j, i)];
        }
    }
    tr
}

/// Soft-thresholding operator `sign(a)·max(|a| − λ, 0)`.
fn soft_threshold(a: f64, lambda: f64) -> f64 {
    if a > lambda {
        a - lambda
    } else if a < -lambda {
        a + lambda
    } else {
        0.0
    }
}

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