regression-diagnostics 0.2.0

use nalgebra::{DMatrix, DVector};
use ndarray::{Array1, Array2, ArrayView1, ArrayView2};

use super::Method;
use crate::error::{RegressionError, Result};
use crate::linalg::{dmatrix_from_rows, dvector_from_slice};
use crate::optimize::nelder_mead;

/// One **random-effect term**: a grouping factor plus the columns whose
/// coefficients vary randomly across its groups.
///
/// * A **random intercept** for a factor is [`RandomEffect::intercept`] (a single
///   column of ones).
/// * A **random slope** (with intercept) passes a design with an intercept column
///   and the slope covariate via [`RandomEffect::new`]; each group then gets its
///   own correlated `(intercept, slope)` pair, `~ N(0, Σ)`.
///
/// Supplying several terms with **different** grouping factors gives a **crossed
/// / nested** model.
#[derive(Debug, Clone)]
pub struct RandomEffect {
    groups: Vec<usize>,
    z: Array2<f64>,
}

impl RandomEffect {
    /// A random **intercept** for the factor whose per-observation labels are
    /// `groups` (arbitrary integers).
    pub fn intercept(groups: &[usize]) -> Self {
        let z = Array2::<f64>::ones((groups.len(), 1));
        Self {
            groups: groups.to_vec(),
            z,
        }
    }

    /// A general random-effect term: `groups` labels and a per-observation design
    /// `z` (`n × k`) whose `k` columns have group-varying coefficients. Include a
    /// column of ones for a random intercept alongside random slopes.
    pub fn new(groups: &[usize], z: Array2<f64>) -> Self {
        Self {
            groups: groups.to_vec(),
            z,
        }
    }
}

/// Internal per-term layout after densifying groups.
struct TermLayout {
    /// Densified group label per observation.
    group_of: Vec<usize>,
    /// Per-observation random-effect design (`n × k`).
    z: Array2<f64>,
    k: usize,
    n_groups: usize,
    /// Column offset of this term's block within the full `Z` / `b`.
    offset: usize,
    /// Offset of this term's parameters within `θ`.
    param_offset: usize,
}

/// A fitted **general linear mixed model** with one or more random-effect terms —
/// random slopes and/or crossed & nested grouping factors — estimated by REML
/// (default) or ML.
///
/// The model is `y = Xβ + Zb + ε`, `b ~ N(0, G)`, `ε ~ N(0, σ²_e I)`, where `Z`
/// and `G` are assembled from the supplied [`RandomEffect`] terms (`G` is
/// block-diagonal, repeating each term's `k × k` covariance across its groups).
/// Estimation profiles `β` (by GLS) and `σ²_e` out analytically and optimizes the
/// remaining **relative covariance** parameters with a Nelder–Mead search, using
/// a dense Cholesky solve of the `n × n` marginal covariance at each step.
///
/// For the single random-intercept case prefer the closed-form
/// [`LinearMixedModel`](super::LinearMixedModel); this type handles everything
/// beyond it, and reduces to it exactly for one intercept term.
///
/// # Scale
///
/// The dense solve is `O(n³)` — appropriate for the grouped datasets these
/// diagnostics target, not for very large `n`.
#[derive(Debug, Clone)]
pub struct MixedModel {
    coefficients: Array1<f64>,
    cov_beta: Array2<f64>,
    var_residual: f64,
    /// Per-term estimated covariance matrices `Σ_term` (`k × k`).
    term_cov: Vec<Array2<f64>>,
    /// Per-term BLUPs, shape `n_groups × k`.
    term_blups: Vec<Array2<f64>>,
    log_likelihood: f64,
    method: Method,
    n: usize,
    p: usize,
    q: usize,
}

impl MixedModel {
    /// Fit by REML. `X` holds the fixed effects (intercept included); `terms` are
    /// the random-effect terms.
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`].
    /// * [`RegressionError::NoResidualDegreesOfFreedom`] if `n ≤ p`.
    /// * [`RegressionError::InvalidResponse`] if no terms are given.
    /// * [`RegressionError::RankDeficient`] if the GLS system is singular.
    pub fn new(x: Array2<f64>, y: Array1<f64>, terms: Vec<RandomEffect>) -> Result<Self> {
        Self::with_method(x, y, terms, Method::Reml)
    }

    /// Like [`MixedModel::new`] with an explicit [`Method`].
    pub fn with_method(
        x: Array2<f64>,
        y: Array1<f64>,
        terms: Vec<RandomEffect>,
        method: Method,
    ) -> 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(),
            });
        }
        if n <= p {
            return Err(RegressionError::NoResidualDegreesOfFreedom {
                n,
                p,
                df: n as isize - p as isize,
            });
        }
        if terms.is_empty() {
            return Err(RegressionError::InvalidResponse {
                msg: "a mixed model needs at least one random-effect term".into(),
            });
        }

        // Build per-term layout and the full Z (n × q).
        let mut layouts = Vec::with_capacity(terms.len());
        let mut q = 0usize;
        let mut n_theta = 0usize;
        for term in &terms {
            if term.groups.len() != n || term.z.nrows() != n {
                return Err(RegressionError::ShapeMismatch {
                    what: "random-effect term length vs X rows",
                    expected: n,
                    got: term.groups.len().min(term.z.nrows()),
                });
            }
            let group_of = densify(&term.groups);
            let n_groups = group_of.iter().copied().max().map_or(0, |m| m + 1);
            let k = term.z.ncols();
            let n_params = k * (k + 1) / 2;
            layouts.push(TermLayout {
                group_of,
                z: term.z.clone(),
                k,
                n_groups,
                offset: q,
                param_offset: n_theta,
            });
            q += n_groups * k;
            n_theta += n_params;
        }

        // Full Z (n × q): observation i contributes term.z[i, c] into the column
        // for (its group, component c) within the term's block.
        let mut z_full = DMatrix::<f64>::zeros(n, q);
        for lay in &layouts {
            for i in 0..n {
                let g = lay.group_of[i];
                for c in 0..lay.k {
                    z_full[(i, lay.offset + g * lay.k + c)] = lay.z[(i, c)];
                }
            }
        }

        let xd = dmatrix_from_rows(n, p, x.as_standard_layout().as_slice().unwrap());
        let yd = dvector_from_slice(y.as_standard_layout().as_slice().unwrap());

        let ctx = Ctx {
            x: &xd,
            y: &yd,
            z: &z_full,
            layouts: &layouts,
            n,
            p,
            q,
            method,
        };

        // Optimize the relative covariance parameters θ. Initialize each term's
        // Cholesky factor to the identity (Δ = I).
        let mut theta0 = vec![0.0; n_theta];
        for lay in &layouts {
            // Diagonal entries of L to 1, off-diagonals 0.
            let mut idx = lay.param_offset;
            for r in 0..lay.k {
                for c in 0..=r {
                    theta0[idx] = if r == c { 1.0 } else { 0.0 };
                    idx += 1;
                }
            }
        }
        let obj = |t: &[f64]| ctx.objective(t).unwrap_or(f64::INFINITY);
        let theta = nelder_mead(obj, &theta0, 0.2, 1e-10, 5000);

        // Final quantities at θ̂.
        let sol = ctx.solve(&theta)?;
        let dof = match method {
            Method::Reml => (n - p) as f64,
            Method::Ml => n as f64,
        };
        let var_residual = sol.rmr / dof;

        let coefficients = Array1::from_shape_fn(p, |j| sol.beta[j]);
        let cov_beta = Array2::from_shape_fn((p, p), |(i, j)| sol.a_inv[(i, j)] * var_residual);

        // Per-term covariance Σ = Δ·σ²_e and BLUPs b̂ = D Zᵀ M⁻¹ r.
        let b = &sol.d * ctx.z.transpose() * &sol.minv_r; // q-vector (relative)
        let mut term_cov = Vec::with_capacity(layouts.len());
        let mut term_blups = Vec::with_capacity(layouts.len());
        for lay in &layouts {
            let delta = relative_covariance(&theta, lay);
            let sigma = Array2::from_shape_fn((lay.k, lay.k), |(i, j)| delta[(i, j)] * var_residual);
            term_cov.push(sigma);
            let blup = Array2::from_shape_fn((lay.n_groups, lay.k), |(g, c)| {
                b[lay.offset + g * lay.k + c]
            });
            term_blups.push(blup);
        }

        let log_likelihood = -0.5 * ctx.objective(&theta)? - ctx.log_const();

        Ok(Self {
            coefficients,
            cov_beta,
            var_residual,
            term_cov,
            term_blups,
            log_likelihood,
            method,
            n,
            p,
            q,
        })
    }

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

    /// Number of fixed-effect coefficients.
    pub fn n_parameters(&self) -> usize {
        self.p
    }

    /// Total number of random-effect coefficients across all terms.
    pub fn n_random_effects(&self) -> usize {
        self.q
    }

    /// Number of random-effect terms.
    pub fn n_terms(&self) -> usize {
        self.term_cov.len()
    }

    /// Estimation method used.
    pub fn method(&self) -> Method {
        self.method
    }

    /// Fixed-effect coefficients `β̂`.
    pub fn coefficients(&self) -> ArrayView1<'_, f64> {
        self.coefficients.view()
    }

    /// Fixed-effect covariance `σ̂²_e (XᵀV⁻¹X)⁻¹`.
    pub fn covariance(&self) -> ArrayView2<'_, f64> {
        self.cov_beta.view()
    }

    /// Fixed-effect standard errors.
    pub fn coefficient_standard_errors(&self) -> Array1<f64> {
        Array1::from_shape_fn(self.p, |j| self.cov_beta[(j, j)].max(0.0).sqrt())
    }

    /// Residual (within-group) variance `σ̂²_e`.
    pub fn residual_variance(&self) -> f64 {
        self.var_residual
    }

    /// Estimated covariance matrix `Σ̂` (`k × k`) of random-effect term `t`; the
    /// diagonal holds the intercept/slope variances, the off-diagonal their
    /// covariance.
    pub fn term_covariance(&self, t: usize) -> ArrayView2<'_, f64> {
        self.term_cov[t].view()
    }

    /// BLUPs for random-effect term `t`, shape `n_groups × k` (row = group,
    /// column = the term's component).
    pub fn random_effects(&self, t: usize) -> ArrayView2<'_, f64> {
        self.term_blups[t].view()
    }

    /// The profile log-likelihood (REML or ML) at the estimate.
    pub fn log_likelihood(&self) -> f64 {
        self.log_likelihood
    }

    /// AIC. Under ML the parameter count is `p` plus the number of covariance
    /// parameters plus one for `σ²_e`; under REML only the covariance parameters
    /// and `σ²_e` are counted.
    pub fn aic(&self) -> f64 {
        let n_cov: usize = self.term_cov.iter().map(|c| {
            let k = c.nrows();
            k * (k + 1) / 2
        }).sum();
        let k = match self.method {
            Method::Ml => self.p as f64 + n_cov as f64 + 1.0,
            Method::Reml => n_cov as f64 + 1.0,
        };
        -2.0 * self.log_likelihood + 2.0 * k
    }
}

/// Cached matrices for the profiled-likelihood search.
struct Ctx<'a> {
    x: &'a DMatrix<f64>,
    y: &'a DVector<f64>,
    z: &'a DMatrix<f64>,
    layouts: &'a [TermLayout],
    n: usize,
    p: usize,
    q: usize,
    method: Method,
}

struct GlsSolve {
    beta: DVector<f64>,
    a_inv: DMatrix<f64>,
    rmr: f64,
    /// `M⁻¹ r` (used to form the BLUPs).
    minv_r: DVector<f64>,
    /// The relative random-effect covariance `D` (`q × q`, block diagonal).
    d: DMatrix<f64>,
}

impl Ctx<'_> {
    /// Build the relative covariance `D`, then GLS-solve at `θ`.
    fn solve(&self, theta: &[f64]) -> Result<GlsSolve> {
        // D (q × q): block diagonal, per term Δ repeated across its groups.
        let mut d = DMatrix::<f64>::zeros(self.q, self.q);
        for lay in self.layouts {
            let delta = relative_covariance(theta, lay);
            for g in 0..lay.n_groups {
                let base = lay.offset + g * lay.k;
                for a in 0..lay.k {
                    for b in 0..lay.k {
                        d[(base + a, base + b)] = delta[(a, b)];
                    }
                }
            }
        }
        // M = I + Z D Zᵀ.
        let mut m = self.z * &d * self.z.transpose();
        for i in 0..self.n {
            m[(i, i)] += 1.0;
        }
        let chol = m.clone().cholesky().ok_or(RegressionError::RankDeficient)?;
        let minv = chol.inverse();

        let xtminv = self.x.transpose() * &minv; // p × n
        let a = &xtminv * self.x; // p × p
        let a_inv = a.try_inverse().ok_or(RegressionError::RankDeficient)?;
        let beta = &a_inv * (&xtminv * self.y);
        let r = self.y - self.x * &beta;
        let minv_r = &minv * &r;
        let rmr = (r.transpose() * &minv_r)[(0, 0)];

        Ok(GlsSolve {
            beta,
            a_inv,
            rmr,
            minv_r,
            d,
        })
    }

    /// Profiled `−2ℓ` (up to the additive constant) to minimize over `θ`.
    fn objective(&self, theta: &[f64]) -> Result<f64> {
        // Recompute M and its Cholesky for the log-determinant.
        let mut d = DMatrix::<f64>::zeros(self.q, self.q);
        for lay in self.layouts {
            let delta = relative_covariance(theta, lay);
            for g in 0..lay.n_groups {
                let base = lay.offset + g * lay.k;
                for a in 0..lay.k {
                    for b in 0..lay.k {
                        d[(base + a, base + b)] = delta[(a, b)];
                    }
                }
            }
        }
        let mut m = self.z * &d * self.z.transpose();
        for i in 0..self.n {
            m[(i, i)] += 1.0;
        }
        let chol = m.cholesky().ok_or(RegressionError::RankDeficient)?;
        let ln_det_m = 2.0 * chol.l().diagonal().iter().map(|v| v.ln()).sum::<f64>();

        let sol = self.solve(theta)?;
        let dof = match self.method {
            Method::Reml => (self.n - self.p) as f64,
            Method::Ml => self.n as f64,
        };
        let sigma2 = (sol.rmr / dof).max(1e-300);
        let mut obj = dof * sigma2.ln() + ln_det_m;
        if self.method == Method::Reml {
            // + ln det(Xᵀ M⁻¹ X) = − ln det(a_inv), since a_inv = (Xᵀ M⁻¹ X)⁻¹.
            let det_ainv = sol.a_inv.determinant();
            obj -= det_ainv.abs().max(1e-300).ln();
        }
        Ok(obj)
    }

    fn log_const(&self) -> f64 {
        let dof = match self.method {
            Method::Reml => (self.n - self.p) as f64,
            Method::Ml => self.n as f64,
        };
        0.5 * dof * ((2.0 * std::f64::consts::PI).ln() + 1.0)
    }
}

/// Assemble a term's relative covariance `Δ = L Lᵀ` from the free lower-triangular
/// parameters in `θ`.
fn relative_covariance(theta: &[f64], lay: &TermLayout) -> DMatrix<f64> {
    let k = lay.k;
    let mut l = DMatrix::<f64>::zeros(k, k);
    let mut idx = lay.param_offset;
    for r in 0..k {
        for c in 0..=r {
            l[(r, c)] = theta[idx];
            idx += 1;
        }
    }
    &l * l.transpose()
}

/// Remap arbitrary integer labels to `0..g`, preserving first-seen order.
fn densify(labels: &[usize]) -> Vec<usize> {
    let mut map = std::collections::BTreeMap::new();
    labels
        .iter()
        .map(|&l| {
            let next = map.len();
            *map.entry(l).or_insert(next)
        })
        .collect()
}