regression-diagnostics 0.2.0

use ndarray::{Array1, Array2, ArrayView1};
use statrs::distribution::{ContinuousCDF, Normal};
use statrs::function::gamma::ln_gamma;

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

/// Conditional response family for a [`GlmmFit`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GlmmFamily {
    /// **Poisson** counts with a log link.
    Poisson,
    /// **Bernoulli** `0/1` responses with a logit link.
    Binomial,
}

impl GlmmFamily {
    fn inverse_link(self, eta: f64) -> f64 {
        match self {
            GlmmFamily::Poisson => eta.exp(),
            GlmmFamily::Binomial => {
                if eta >= 0.0 {
                    1.0 / (1.0 + (-eta).exp())
                } else {
                    let e = eta.exp();
                    e / (1.0 + e)
                }
            }
        }
    }

    /// Per-observation conditional log-likelihood `ℓ(yᵢ | ηᵢ)`.
    fn loglik(self, y: f64, eta: f64) -> f64 {
        match self {
            GlmmFamily::Poisson => y * eta - eta.exp() - ln_gamma(y + 1.0),
            GlmmFamily::Binomial => {
                // y·η − ln(1 + e^η), stable.
                let lse = if eta > 0.0 {
                    eta + (-eta).exp().ln_1p()
                } else {
                    eta.exp().ln_1p()
                };
                y * eta - lse
            }
        }
    }

    /// GLM weight `wᵢ = −∂²ℓ/∂η² = V(μ)` at the canonical link (μ for Poisson,
    /// μ(1−μ) for Bernoulli).
    fn weight(self, mu: f64) -> f64 {
        match self {
            GlmmFamily::Poisson => mu,
            GlmmFamily::Binomial => mu * (1.0 - mu),
        }
    }

    fn validate(self, y: &Array1<f64>) -> Result<()> {
        match self {
            GlmmFamily::Poisson => {
                for &v in y.iter() {
                    if !v.is_finite() || v < 0.0 {
                        return Err(RegressionError::InvalidResponse {
                            msg: format!("Poisson GLMM response must be a non-negative count, found {v}"),
                        });
                    }
                }
            }
            GlmmFamily::Binomial => {
                for &v in y.iter() {
                    if v != 0.0 && v != 1.0 {
                        return Err(RegressionError::InvalidResponse {
                            msg: format!("binomial GLMM response must be 0 or 1, found {v}"),
                        });
                    }
                }
            }
        }
        Ok(())
    }
}

/// A fitted **random-intercept generalized linear mixed model** (GLMM),
///
/// `g(E[yᵢⱼ | bⱼ]) = xᵢⱼᵀβ + bⱼ`,  `bⱼ ~ N(0, σ_b²)`,
///
/// for a non-Gaussian conditional [`GlmmFamily`] (Poisson counts or Bernoulli
/// outcomes) — the generalized counterpart to
/// [`LinearMixedModel`](super::LinearMixedModel). The intractable integral over
/// the random effects is handled by the **Laplace approximation**: an inner
/// Newton loop finds each group's conditional mode, and an outer Nelder–Mead
/// search maximizes the resulting approximate marginal likelihood over the fixed
/// effects `β` and the random-intercept standard deviation `σ_b`.
///
/// Fixed-effect standard errors come from the numerical observed information of
/// the Laplace log-likelihood. The random-intercept variance is reported as
/// `σ_b`; as it shrinks toward zero the model approaches the corresponding plain
/// GLM.
#[derive(Debug, Clone)]
pub struct GlmmFit {
    family: GlmmFamily,
    coefficients: Array1<f64>,
    sigma_b: f64,
    cov_beta: Array2<f64>,
    blups: Array1<f64>,
    log_likelihood: f64,
    n: usize,
    p: usize,
    n_groups: usize,
}

impl GlmmFit {
    /// Fit a random-intercept GLMM of `y` on fixed-effect design `X` with group
    /// labels `groups`, under `family`, by Laplace-approximate maximum
    /// likelihood.
    ///
    /// `X` carries the fixed effects including an intercept.
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`].
    /// * [`RegressionError::InvalidResponse`] for out-of-support responses or
    ///   fewer than two groups.
    /// * [`RegressionError::NotConverged`] if the optimizer fails to find a finite
    ///   optimum.
    pub fn new(
        x: Array2<f64>,
        y: Array1<f64>,
        groups: &[usize],
        family: GlmmFamily,
    ) -> 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 || groups.len() != n {
            return Err(RegressionError::ShapeMismatch {
                what: "y/groups length vs X rows",
                expected: n,
                got: y.len().min(groups.len()),
            });
        }
        family.validate(&y)?;

        // Densify groups and collect per-group row indices.
        let mut map = std::collections::BTreeMap::new();
        for &g in groups {
            let next = map.len();
            map.entry(g).or_insert(next);
        }
        let n_groups = map.len();
        if n_groups < 2 {
            return Err(RegressionError::InvalidResponse {
                msg: "a GLMM needs at least two groups".into(),
            });
        }
        let mut group_rows: Vec<Vec<usize>> = vec![Vec::new(); n_groups];
        for (i, &lab) in groups.iter().enumerate() {
            group_rows[map[&lab]].push(i);
        }

        // Parameter vector θ = (β, ln σ_b). Negative Laplace log-likelihood.
        let neg_ll = |theta: &[f64]| -> f64 {
            let sigma_b = theta[p].exp();
            if !sigma_b.is_finite() || sigma_b <= 0.0 {
                return f64::INFINITY;
            }
            match laplace_loglik(&x, &y, &group_rows, theta, sigma_b, family) {
                Some((ll, _)) if ll.is_finite() => -ll,
                _ => f64::INFINITY,
            }
        };

        // Warm start: intercept-only mean on the link scale, β slopes 0, σ_b = 0.5.
        let mut theta0 = vec![0.0; p + 1];
        let ybar = y.sum() / n as f64;
        theta0[0] = match family {
            GlmmFamily::Poisson => ybar.max(1e-3).ln(),
            GlmmFamily::Binomial => (ybar.clamp(1e-3, 1.0 - 1e-3)
                / (1.0 - ybar.clamp(1e-3, 1.0 - 1e-3)))
            .ln(),
        };
        theta0[p] = 0.5_f64.ln();

        let neg_ll_ref = &neg_ll;
        let theta = nelder_mead(neg_ll_ref, &theta0, 0.2, 1e-9, 8000);
        let sigma_b = theta[p].exp();
        let (ll, blups_vec) =
            laplace_loglik(&x, &y, &group_rows, &theta, sigma_b, family).ok_or(
                RegressionError::NotConverged {
                    iterations: 8000,
                    msg: "GLMM Laplace optimizer failed to find a finite optimum".into(),
                },
            )?;

        // Fixed-effect covariance from the numerical observed information.
        let grad = |t: &[f64]| -> Vec<f64> {
            let mut g = vec![0.0; p + 1];
            for j in 0..=p {
                let h = 1e-5 * t[j].abs().max(1.0);
                let mut tp = t.to_vec();
                let mut tm = t.to_vec();
                tp[j] += h;
                tm[j] -= h;
                g[j] = (neg_ll(&tp) - neg_ll(&tm)) / (2.0 * h);
            }
            g
        };
        let hess = numerical_hessian(grad, &theta);
        let flat: Vec<f64> = hess.iter().flat_map(|r| r.iter().copied()).collect();
        let cov_beta = match dmatrix_from_rows(p + 1, p + 1, &flat).try_inverse() {
            Some(inv) => Array2::from_shape_fn((p, p), |(i, j)| inv[(i, j)]),
            None => Array2::from_elem((p, p), f64::NAN),
        };

        let coefficients = Array1::from_shape_fn(p, |j| theta[j]);
        let blups = Array1::from(blups_vec);

        Ok(Self {
            family,
            coefficients,
            sigma_b,
            cov_beta,
            blups,
            log_likelihood: ll,
            n,
            p,
            n_groups,
        })
    }

    /// The conditional family.
    pub fn family(&self) -> GlmmFamily {
        self.family
    }

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

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

    /// Number of groups.
    pub fn n_groups(&self) -> usize {
        self.n_groups
    }

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

    /// Random-intercept standard deviation `σ̂_b`.
    pub fn sigma_b(&self) -> f64 {
        self.sigma_b
    }

    /// Random-intercept variance `σ̂_b²`.
    pub fn group_variance(&self) -> f64 {
        self.sigma_b * self.sigma_b
    }

    /// Fixed-effect covariance (inverse observed information of the Laplace
    /// likelihood).
    pub fn covariance(&self) -> ndarray::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())
    }

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

    /// Predicted (conditional-mode) random intercepts `b̂_j`, in densified group
    /// order.
    pub fn random_effects(&self) -> ArrayView1<'_, f64> {
        self.blups.view()
    }

    /// The Laplace-approximate log-likelihood at the estimate.
    pub fn log_likelihood(&self) -> f64 {
        self.log_likelihood
    }

    /// AIC, `−2ℓ + 2(p + 1)` (fixed effects plus the variance component).
    pub fn aic(&self) -> f64 {
        -2.0 * self.log_likelihood + 2.0 * (self.p as f64 + 1.0)
    }
}

/// Laplace log-likelihood at `θ = (β, ln σ_b)` with the given `sigma_b`, plus the
/// conditional modes `û_j`. Returns `None` if the inner Newton diverges.
fn laplace_loglik(
    x: &Array2<f64>,
    y: &Array1<f64>,
    group_rows: &[Vec<usize>],
    theta: &[f64],
    sigma_b: f64,
    family: GlmmFamily,
) -> Option<(f64, Vec<f64>)> {
    let p = x.ncols();
    let s2 = sigma_b * sigma_b;
    let mut total = 0.0;
    let mut modes = Vec::with_capacity(group_rows.len());

    for rows in group_rows {
        // Fixed part ηˣ_i = x_iᵀβ.
        let eta_fixed: Vec<f64> = rows
            .iter()
            .map(|&i| (0..p).map(|j| x[(i, j)] * theta[j]).sum::<f64>())
            .collect();

        // Inner Newton for the conditional mode u.
        let mut u = 0.0;
        for _ in 0..100 {
            let mut grad = -u / s2;
            let mut ws = 0.0;
            for (k, &i) in rows.iter().enumerate() {
                let eta = eta_fixed[k] + u;
                let mu = family.inverse_link(eta);
                grad += y[i] - mu;
                ws += family.weight(mu);
            }
            let h = ws + 1.0 / s2;
            let step = grad / h;
            u += step;
            if !u.is_finite() {
                return None;
            }
            if step.abs() < 1e-12 {
                break;
            }
        }

        // Q_j(û) = Σ loglik_i(û) − û²/(2σ_b²); evaluate w_sum at û.
        let mut q = -u * u / (2.0 * s2);
        let mut w_sum = 0.0;
        for (k, &i) in rows.iter().enumerate() {
            let eta = eta_fixed[k] + u;
            q += family.loglik(y[i], eta);
            w_sum += family.weight(family.inverse_link(eta));
        }
        // ln L_j ≈ Q_j(û) − ½ ln(1 + σ_b² W_j).
        let contrib = q - 0.5 * (1.0 + s2 * w_sum).ln();
        if !contrib.is_finite() {
            return None;
        }
        total += contrib;
        modes.push(u);
    }
    Some((total, modes))
}