regression-diagnostics 0.2.0

use ndarray::{Array1, Array2};

use super::{GlmFit, NegativeBinomial};
use crate::error::{RegressionError, Result};

/// Lower/upper bounds for the dispersion search, on the log scale. Above the
/// upper bound the negative binomial is indistinguishable from Poisson.
const THETA_MIN: f64 = 1e-3;
const THETA_MAX: f64 = 1e6;

/// Fit a negative-binomial GLM **estimating the dispersion `θ` jointly** with the
/// coefficients, by maximizing the profile log-likelihood
/// `θ ↦ ℓ(β̂(θ), θ)` over `θ` (a one-dimensional golden-section search, refitting
/// the coefficients at each candidate `θ`).
///
/// Returns the fitted [`GlmFit<NegativeBinomial>`] at the estimated `θ̂`; read the
/// value back with `fit.family().theta()`. Use this when the count data are
/// over-dispersed and you do not have `θ` in advance — the counterpart to
/// [`GlmFit::new`] with a fixed [`NegativeBinomial`].
///
/// # Boundary behavior
///
/// If the data show no over-dispersion the profile keeps rising as `θ → ∞`
/// (approaching Poisson); the search then returns the capped `θ̂ ≈ 1e6`, which is
/// the signal to prefer a plain [`Poisson`](super::Poisson) fit.
///
/// # Errors
///
/// Propagates the [`GlmFit`] fitting errors (empty/mismatched input, invalid
/// response, non-convergence).
pub fn fit_negative_binomial(x: Array2<f64>, y: Array1<f64>) -> Result<GlmFit<NegativeBinomial>> {
    // Objective: negative profile log-likelihood at dispersion exp(t).
    let neg_ll = |t: f64| -> f64 {
        let theta = t.exp();
        match NegativeBinomial::new(theta)
            .and_then(|fam| GlmFit::new(fam, x.clone(), y.clone()))
        {
            Ok(fit) => -fit.log_likelihood(),
            Err(_) => f64::INFINITY,
        }
    };

    // Golden-section minimization of neg_ll over t = ln θ ∈ [ln THETA_MIN, ln THETA_MAX].
    let phi = (5.0_f64.sqrt() - 1.0) / 2.0;
    let mut lo = THETA_MIN.ln();
    let mut hi = THETA_MAX.ln();
    let mut c = hi - phi * (hi - lo);
    let mut d = lo + phi * (hi - lo);
    let mut fc = neg_ll(c);
    let mut fd = neg_ll(d);
    for _ in 0..200 {
        if fc < fd {
            hi = d;
            d = c;
            fd = fc;
            c = hi - phi * (hi - lo);
            fc = neg_ll(c);
        } else {
            lo = c;
            c = d;
            fc = fd;
            d = lo + phi * (hi - lo);
            fd = neg_ll(d);
        }
        if (hi - lo) < 1e-8 {
            break;
        }
    }
    let t_hat = 0.5 * (lo + hi);
    let theta_hat = t_hat.exp().clamp(THETA_MIN, THETA_MAX);

    let fam = NegativeBinomial::new(theta_hat).map_err(|_| RegressionError::InvalidParameter {
        msg: "estimated θ out of range".into(),
    })?;
    GlmFit::new(fam, x, y)
}