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;

/// A fitted **proportional-odds (ordinal logistic)** model for an ordered
/// categorical response with `K ≥ 2` levels `0 < 1 < … < K−1`.
///
/// Uses the cumulative-logit parametrization of R's `MASS::polr`:
///
/// `logit P(yᵢ ≤ k) = αₖ − xᵢᵀβ`,  `k = 0 … K−2`,
///
/// with strictly increasing thresholds `α₀ < α₁ < … < α_{K−2}` and a **single**
/// coefficient vector `β` shared across all thresholds (the proportional-odds
/// assumption). Positive `βⱼ` raises the odds of falling in a *higher* category.
/// The design matrix carries **no intercept** — the thresholds play that role.
///
/// Fit by Newton–Raphson on the full `(K−1) + p` parameter vector (thresholds
/// then coefficients) with a backtracking line search. With `K = 2` this reduces
/// exactly to binary [`LogisticFit`](crate::logistic::LogisticFit): `β` equals
/// the logistic slopes and `α₀` equals the negated logistic intercept.
#[derive(Debug, Clone)]
pub struct OrdinalFit {
    x: Array2<f64>,
    y: Array1<f64>,
    /// Threshold (cutpoint) parameters `α₀ < … < α_{K−2}`, length `K−1`.
    thresholds: Array1<f64>,
    /// Shared coefficient vector `β`, length `p`.
    coefficients: Array1<f64>,
    /// Fitted class probabilities, shape `n × K`.
    probabilities: Array2<f64>,
    /// Covariance of the stacked `(K−1)+p` parameter vector (thresholds first).
    cov: Array2<f64>,
    log_likelihood: f64,
    iterations: usize,
    n: usize,
    p: usize,
    k: usize,
}

impl OrdinalFit {
    /// Fit a proportional-odds model of ordered `y` on `X` (default: up to 100
    /// Newton iterations, tolerance `1e-10`).
    ///
    /// The response holds integer levels `0 … K−1`, every level present; `K` is
    /// inferred as `max(y) + 1`. **Do not include an intercept column** in `X`.
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`].
    /// * [`RegressionError::InvalidResponse`] for non-integer/negative labels, a
    ///   missing level, or fewer than two levels.
    /// * [`RegressionError::RankDeficient`] if the information matrix is singular.
    /// * [`RegressionError::NotConverged`] if Newton's method fails to converge.
    pub fn new(x: Array2<f64>, y: Array1<f64>) -> Result<Self> {
        Self::with_options(x, y, 100, 1e-10)
    }

    /// Like [`OrdinalFit::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(),
            });
        }
        let k = validate_labels(&y)?;
        let n_thresh = k - 1;
        let m = n_thresh + p;

        // Initialize thresholds from cumulative class frequencies, β = 0.
        let mut counts = vec![0.0_f64; k];
        for &yi in y.iter() {
            counts[yi as usize] += 1.0;
        }
        let mut theta = Array1::<f64>::zeros(m);
        let mut cum = 0.0;
        for kk in 0..n_thresh {
            cum += counts[kk];
            let prop = (cum / n as f64).clamp(1e-4, 1.0 - 1e-4);
            theta[kk] = (prop / (1.0 - prop)).ln();
        }

        let mut iterations = 0usize;
        let mut converged = false;
        let mut cov = Array2::<f64>::zeros((m, m));

        let mut nll = neg_log_likelihood(&x, &y, &theta, k);
        while iterations < max_iter {
            iterations += 1;

            let grad = gradient(&x, &y, &theta, k);
            let hess = hessian(&x, &y, &theta, k);

            let hess_dm = dmatrix_from_rows(m, m, hess.as_standard_layout().as_slice().unwrap());
            let inv = hess_dm.try_inverse().ok_or(RegressionError::RankDeficient)?;
            let inv_arr = Array2::from_shape_fn((m, m), |(i, j)| inv[(i, j)]);
            cov = inv_arr.clone();

            // Newton step for minimizing the nll: Δ = −H⁻¹ g.
            let mut step = inv_arr.dot(&grad);
            step.mapv_inplace(|v| -v);

            // Backtracking line search: accept the step only if it decreases the
            // nll (and keeps thresholds ordered), else halve.
            let mut scale = 1.0_f64;
            let mut new_theta = &theta + &step;
            let mut new_nll = f64::INFINITY;
            for _ in 0..30 {
                new_theta = &theta + &(&step * scale);
                if thresholds_ordered(&new_theta, n_thresh) {
                    new_nll = neg_log_likelihood(&x, &y, &new_theta, k);
                    if new_nll.is_finite() && new_nll <= nll + 1e-12 {
                        break;
                    }
                }
                scale *= 0.5;
            }

            let max_step = step
                .iter()
                .map(|v| (v * scale).abs())
                .fold(0.0_f64, f64::max);
            theta = new_theta;
            nll = new_nll;

            if !theta.iter().all(|v| v.is_finite()) {
                return Err(RegressionError::NotConverged {
                    iterations,
                    msg: "parameters diverging".into(),
                });
            }
            if max_step < tol {
                converged = true;
                break;
            }
        }

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

        let thresholds = Array1::from_shape_fn(n_thresh, |i| theta[i]);
        let coefficients = Array1::from_shape_fn(p, |i| theta[n_thresh + i]);

        let mut probabilities = Array2::<f64>::zeros((n, k));
        fill_probabilities(&x, &thresholds, &coefficients, &mut probabilities);
        let log_likelihood = -nll;

        Ok(Self {
            x,
            y,
            thresholds,
            coefficients,
            probabilities,
            cov,
            log_likelihood,
            iterations,
            n,
            p,
            k,
        })
    }

    /// Number of observations.
    pub fn n_observations(&self) -> usize {
        self.n
    }

    /// Number of design columns `p`.
    pub fn n_features(&self) -> usize {
        self.p
    }

    /// Number of ordered levels `K`.
    pub fn n_classes(&self) -> usize {
        self.k
    }

    /// Total number of parameters, `(K − 1) + p` (thresholds plus coefficients).
    pub fn n_parameters(&self) -> usize {
        (self.k - 1) + self.p
    }

    /// Newton 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 integer level labels.
    pub fn response(&self) -> ArrayView1<'_, f64> {
        self.y.view()
    }

    /// Threshold (cutpoint) parameters `α₀ < … < α_{K−2}`.
    pub fn thresholds(&self) -> ArrayView1<'_, f64> {
        self.thresholds.view()
    }

    /// Shared coefficient vector `β` (proportional-odds effects).
    pub fn coefficients(&self) -> ArrayView1<'_, f64> {
        self.coefficients.view()
    }

    /// Fitted class probabilities, shape `n × K`.
    pub fn fitted_probabilities(&self) -> ArrayView2<'_, f64> {
        self.probabilities.view()
    }

    /// Covariance of the stacked `(K−1)+p` parameter vector (thresholds first).
    pub fn covariance(&self) -> ArrayView2<'_, f64> {
        self.cov.view()
    }

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

    /// Standard errors of the coefficients `β` (the last `p` diagonal entries of
    /// the covariance).
    pub fn coefficient_standard_errors(&self) -> Array1<f64> {
        let off = self.k - 1;
        Array1::from_shape_fn(self.p, |j| self.cov[(off + j, off + j)].max(0.0).sqrt())
    }

    /// Standard errors of the thresholds `α`.
    pub fn threshold_standard_errors(&self) -> Array1<f64> {
        Array1::from_shape_fn(self.k - 1, |j| self.cov[(j, j)].max(0.0).sqrt())
    }

    /// Wald `z`-statistics for the coefficients, `βⱼ / 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 for the coefficients 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
            }
        })
    }

    /// Residual deviance `−2ℓ`.
    pub fn residual_deviance(&self) -> f64 {
        -2.0 * self.log_likelihood
    }

    /// Deviance of the intercept-only (threshold-only, `β = 0`) model.
    pub fn null_deviance(&self) -> f64 {
        -2.0 * self.null_log_likelihood()
    }

    fn null_log_likelihood(&self) -> f64 {
        let n = self.n as f64;
        let mut counts = vec![0.0_f64; self.k];
        for &yi in self.y.iter() {
            counts[yi as usize] += 1.0;
        }
        counts
            .iter()
            .filter(|&&c| c > 0.0)
            .map(|&c| c * (c / n).ln())
            .sum()
    }

    /// McFadden's pseudo-R², `1 − ℓ/ℓ₀`.
    pub fn mcfadden_r2(&self) -> f64 {
        let ll0 = self.null_log_likelihood();
        if ll0 != 0.0 {
            1.0 - self.log_likelihood / ll0
        } else {
            f64::NAN
        }
    }

    /// Akaike information criterion, `−2ℓ + 2·[(K−1)+p]`.
    pub fn aic(&self) -> f64 {
        self.residual_deviance() + 2.0 * self.n_parameters() as f64
    }

    /// Bayesian information criterion, `−2ℓ + ln(n)·[(K−1)+p]`.
    pub fn bic(&self) -> f64 {
        self.residual_deviance() + (self.n as f64).ln() * self.n_parameters() as f64
    }

    /// Predicted class probabilities for a new design matrix `x`, shape
    /// `rows × K`.
    pub fn predict_proba(&self, x: ArrayView2<'_, f64>) -> Array2<f64> {
        let xo = x.to_owned();
        let mut out = Array2::<f64>::zeros((xo.nrows(), self.k));
        fill_probabilities(&xo, &self.thresholds, &self.coefficients, &mut out);
        out
    }
}

/// Per-observation deviance residuals `√(−2 ln P(yᵢ = cᵢ)) ≥ 0`; they square to
/// the residual deviance.
pub fn deviance_residuals(fit: &OrdinalFit) -> Array1<f64> {
    let y = fit.response();
    let p = fit.fitted_probabilities();
    Array1::from_shape_fn(fit.n_observations(), |i| {
        let pi = p[(i, y[i] as usize)].max(1e-12);
        (-2.0 * pi.ln()).max(0.0).sqrt()
    })
}

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

/// Cumulative CDFs for observation `i`: returns `(P(y≤c), P(y≤c−1))` and their
/// densities `(σ'(A), σ'(B))`, handling the open ends `c = 0` and `c = K−1`.
fn cell_terms(eta: f64, alpha: &[f64], c: usize, k: usize) -> (f64, f64, f64, f64) {
    let (s_a, sp_a) = if c == k - 1 {
        (1.0, 0.0)
    } else {
        let a = alpha[c] - eta;
        let s = sigmoid(a);
        (s, s * (1.0 - s))
    };
    let (s_b, sp_b) = if c == 0 {
        (0.0, 0.0)
    } else {
        let b = alpha[c - 1] - eta;
        let s = sigmoid(b);
        (s, s * (1.0 - s))
    };
    (s_a, s_b, sp_a, sp_b)
}

fn neg_log_likelihood(x: &Array2<f64>, y: &Array1<f64>, theta: &Array1<f64>, k: usize) -> f64 {
    let n = x.nrows();
    let p = x.ncols();
    let n_thresh = k - 1;
    let alpha = &theta.as_slice().unwrap()[0..n_thresh];
    let beta = &theta.as_slice().unwrap()[n_thresh..];
    let mut nll = 0.0;
    for i in 0..n {
        let mut eta = 0.0;
        for j in 0..p {
            eta += x[(i, j)] * beta[j];
        }
        let c = y[i] as usize;
        let (s_a, s_b, _, _) = cell_terms(eta, alpha, c, k);
        let prob = (s_a - s_b).max(1e-12);
        nll -= prob.ln();
    }
    nll
}

fn gradient(x: &Array2<f64>, y: &Array1<f64>, theta: &Array1<f64>, k: usize) -> Array1<f64> {
    let n = x.nrows();
    let p = x.ncols();
    let n_thresh = k - 1;
    let m = n_thresh + p;
    let alpha = &theta.as_slice().unwrap()[0..n_thresh];
    let beta = &theta.as_slice().unwrap()[n_thresh..];
    let mut g = Array1::<f64>::zeros(m); // gradient of the nll
    for i in 0..n {
        let mut eta = 0.0;
        for j in 0..p {
            eta += x[(i, j)] * beta[j];
        }
        let c = y[i] as usize;
        let (s_a, s_b, sp_a, sp_b) = cell_terms(eta, alpha, c, k);
        let prob = (s_a - s_b).max(1e-12);
        // ∂nll/∂α_m = −(1/P)(σ'(A)[m==c] − σ'(B)[m==c−1]).
        if c < n_thresh {
            g[c] -= sp_a / prob;
        }
        if c >= 1 {
            g[c - 1] -= -sp_b / prob;
        }
        // ∂nll/∂β_j = (x_ij/P)(σ'(A) − σ'(B)).
        let common = (sp_a - sp_b) / prob;
        for j in 0..p {
            g[n_thresh + j] += x[(i, j)] * common;
        }
    }
    g
}

/// Observed information (Hessian of the nll) by central differences of the
/// analytic gradient — robust and, for these small parameter vectors, cheap.
fn hessian(x: &Array2<f64>, y: &Array1<f64>, theta: &Array1<f64>, k: usize) -> Array2<f64> {
    let m = theta.len();
    let mut h = Array2::<f64>::zeros((m, m));
    let eps = 1e-6;
    for j in 0..m {
        let mut tp = theta.clone();
        let mut tm = theta.clone();
        let step = eps * theta[j].abs().max(1.0);
        tp[j] += step;
        tm[j] -= step;
        let gp = gradient(x, y, &tp, k);
        let gm = gradient(x, y, &tm, k);
        for i in 0..m {
            h[(i, j)] = (gp[i] - gm[i]) / (2.0 * step);
        }
    }
    // Symmetrize to counter finite-difference asymmetry.
    for i in 0..m {
        for j in (i + 1)..m {
            let avg = 0.5 * (h[(i, j)] + h[(j, i)]);
            h[(i, j)] = avg;
            h[(j, i)] = avg;
        }
    }
    h
}

fn thresholds_ordered(theta: &Array1<f64>, n_thresh: usize) -> bool {
    for k in 1..n_thresh {
        if theta[k] <= theta[k - 1] {
            return false;
        }
    }
    true
}

fn fill_probabilities(
    x: &Array2<f64>,
    alpha: &Array1<f64>,
    beta: &Array1<f64>,
    probs: &mut Array2<f64>,
) {
    let n = x.nrows();
    let p = x.ncols();
    let k = alpha.len() + 1;
    for i in 0..n {
        let mut eta = 0.0;
        for j in 0..p {
            eta += x[(i, j)] * beta[j];
        }
        let mut prev = 0.0;
        for c in 0..k {
            let cdf = if c == k - 1 {
                1.0
            } else {
                sigmoid(alpha[c] - eta)
            };
            probs[(i, c)] = (cdf - prev).max(0.0);
            prev = cdf;
        }
    }
}

fn validate_labels(y: &Array1<f64>) -> Result<usize> {
    let mut max_label = 0usize;
    for &v in y.iter() {
        if !v.is_finite() || v < 0.0 || v.fract() != 0.0 {
            return Err(RegressionError::InvalidResponse {
                msg: format!("ordinal labels must be non-negative integers, found {v}"),
            });
        }
        max_label = max_label.max(v as usize);
    }
    let k = max_label + 1;
    if k < 2 {
        return Err(RegressionError::InvalidResponse {
            msg: "ordinal response needs at least two levels".into(),
        });
    }
    let mut present = vec![false; k];
    for &v in y.iter() {
        present[v as usize] = true;
    }
    if let Some(missing) = present.iter().position(|&b| !b) {
        return Err(RegressionError::InvalidResponse {
            msg: format!("level {missing} has no observations; labels must be 0..K-1 with all present"),
        });
    }
    Ok(k)
}