regression-diagnostics 0.2.0

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

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

/// A fitted lasso-regression model and its diagnostics.
///
/// Lasso solves `min (1/2n)‖y − Xβ‖² + λ‖β_pen‖₁`. Unlike ridge there is no
/// closed form, so this fits by **cyclic coordinate descent** with
/// soft-thresholding — the standard, well-conditioned algorithm (glmnet-style).
///
/// # Intercept and scaling
///
/// A detected constant column is an **unpenalized intercept**. Predictors are
/// **standardized internally** (centered and scaled to unit variance) before the
/// penalty is applied, then coefficients are transformed back to the original
/// scale; the intercept is recovered from the means. Because standardization is
/// internal, `λ` is on the standardized `(1/2n)`-objective scale — not comparable
/// to ridge's `λ`.
///
/// # The natural diagnostic: the active set
///
/// Lasso's defining behavior is that it drives coefficients *exactly* to zero.
/// The size of the surviving [`active_set`](Self::active_set) is an unbiased
/// estimate of the model's degrees of freedom (Zou, Hastie & Tibshirani, 2007),
/// which is what the information criteria here use as the parameter count.
#[derive(Debug, Clone)]
pub struct LassoFit {
    x: Array2<f64>,
    y: Array1<f64>,
    lambda: f64,
    coefficients: Array1<f64>,
    fitted: Array1<f64>,
    residuals: Array1<f64>,
    rss: f64,
    intercept_col: Option<usize>,
    n_nonzero: usize,
    iterations: usize,
    n: usize,
    p: usize,
}

impl LassoFit {
    /// Fit lasso regression of `y` on `X` with penalty `lambda ≥ 0` using
    /// coordinate descent (default tolerance `1e-7`, up to `10_000` sweeps).
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`].
    /// * [`RegressionError::InvalidParameter`] if `lambda < 0`.
    /// * [`RegressionError::NotConverged`] if coordinate descent does not
    ///   converge within the iteration budget.
    pub fn new(x: Array2<f64>, y: Array1<f64>, lambda: f64) -> Result<Self> {
        Self::with_options(x, y, lambda, 1e-7, 10_000)
    }

    /// Like [`LassoFit::new`] but with an explicit convergence tolerance and
    /// maximum number of coordinate-descent sweeps.
    pub fn with_options(
        x: Array2<f64>,
        y: Array1<f64>,
        lambda: 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!("lasso lambda must be >= 0, got {lambda}"),
            });
        }

        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: z_j = (x_j − mean_j) / sd_j, with population sd.
        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 };
        }
        // z: n × q standardized design.
        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 on standardized coefficients β (length q).
        let mut beta = vec![0.0f64; q];
        // residual r = yc − Z β (starts at yc since β = 0).
        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 {
                // ρ_k = (1/n) z_kᵀ r + β_k   (since (1/n) z_kᵀ z_k = 1)
                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, lambda);
                let delta = new - beta[k];
                if delta != 0.0 {
                    // Update residual: r -= z_k * delta
                    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: "lasso coordinate descent did not reach tolerance".into(),
            });
        }

        // Transform back to original scale: β_orig_j = β_std / sd_j.
        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();
        let n_nonzero = slopes_orig.iter().filter(|b| b.abs() > 1e-12).count();

        Ok(Self {
            x,
            y,
            lambda,
            coefficients,
            fitted,
            residuals,
            rss,
            intercept_col,
            n_nonzero,
            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()
    }

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

    /// Lasso coefficients, aligned to the design columns. Penalized coefficients
    /// that were 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 (surviving) 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 — the active-set size, which is
    /// lasso's degrees-of-freedom estimate (Zou–Hastie–Tibshirani).
    pub fn n_nonzero(&self) -> usize {
        self.n_nonzero
    }

    /// Effective degrees of freedom: the active-set size plus one for the
    /// intercept if present.
    pub fn effective_df(&self) -> f64 {
        self.n_nonzero as f64 + if self.has_intercept() { 1.0 } else { 0.0 }
    }

    /// 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 active-set-based 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 active-set-based degrees of freedom.
    pub fn bic(&self) -> f64 {
        -2.0 * self.log_likelihood() + (self.n as f64).ln() * self.effective_df()
    }
}

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