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, dvector_from_slice};
use crate::optimize::{nelder_mead, numerical_hessian};

/// Error / baseline distribution of an [`AftFit`], on the log-time scale
/// `log T = xᵀβ + σ·W`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AftDistribution {
    /// `W` standard extreme-value (Gumbel-min): `T` is **Weibull**. The workhorse
    /// AFT model; also proportional-hazards.
    Weibull,
    /// Weibull with the scale fixed at `σ = 1`: `T` is **exponential** (constant
    /// hazard).
    Exponential,
    /// `W` standard normal: `T` is **log-normal**.
    LogNormal,
    /// `W` standard logistic: `T` is **log-logistic**.
    LogLogistic,
}

impl AftDistribution {
    /// Log standard density `ln f₀(w)`.
    fn log_pdf(self, w: f64) -> f64 {
        match self {
            AftDistribution::Weibull | AftDistribution::Exponential => w - w.exp(),
            AftDistribution::LogNormal => {
                -0.5 * (2.0 * std::f64::consts::PI).ln() - 0.5 * w * w
            }
            AftDistribution::LogLogistic => {
                // w − 2 ln(1 + e^w), computed stably.
                w - 2.0 * log1p_exp(w)
            }
        }
    }

    /// Log standard survivor `ln S₀(w)`.
    fn log_surv(self, w: f64) -> f64 {
        match self {
            AftDistribution::Weibull | AftDistribution::Exponential => -w.exp(),
            AftDistribution::LogNormal => {
                let n = Normal::new(0.0, 1.0).unwrap();
                (1.0 - n.cdf(w)).max(1e-300).ln()
            }
            AftDistribution::LogLogistic => -log1p_exp(w),
        }
    }

    /// Median of the standard variate `W` (for median survival predictions).
    fn median_w(self) -> f64 {
        match self {
            // median of extreme-value-min: ln(ln 2).
            AftDistribution::Weibull | AftDistribution::Exponential => (2.0_f64.ln()).ln(),
            AftDistribution::LogNormal | AftDistribution::LogLogistic => 0.0,
        }
    }

    fn scale_fixed(self) -> bool {
        matches!(self, AftDistribution::Exponential)
    }
}

/// Numerically stable `ln(1 + eˣ)`.
fn log1p_exp(x: f64) -> f64 {
    if x > 0.0 {
        x + (-x).exp().ln_1p()
    } else {
        x.exp().ln_1p()
    }
}

/// A fitted **accelerated failure time (AFT)** parametric survival model,
///
/// `log Tᵢ = xᵢᵀβ + σ·Wᵢ`,
///
/// where `W` has the standard [`AftDistribution`]. Unlike the semiparametric Cox
/// model this specifies the full baseline, so it yields absolute time
/// predictions (e.g. median survival) and coefficients read as **log
/// time-acceleration**: `exp(βⱼ) > 1` multiplies survival time. Fit by maximum
/// likelihood over right-censored data (Nelder–Mead on the log-likelihood, with
/// standard errors from the numerical observed information).
#[derive(Debug, Clone)]
pub struct AftFit {
    dist: AftDistribution,
    coefficients: Array1<f64>,
    scale: f64,
    cov: Array2<f64>,
    log_likelihood: f64,
    n: usize,
    p: usize,
}

impl AftFit {
    /// Fit an AFT model of survival `(time, event)` on `X` under `dist`.
    ///
    /// `event[i]` is `1.0` for an observed event, `0.0` for right-censoring. `X`
    /// carries the fixed effects **including an intercept** (the model has a
    /// location term). Returns the standard-error covariance for the `β`
    /// coefficients (the scale is reported separately).
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] / [`RegressionError::ShapeMismatch`].
    /// * [`RegressionError::InvalidResponse`] for non-positive times or event
    ///   flags outside `{0, 1}`, or no events.
    /// * [`RegressionError::NotConverged`] if the optimizer stalls.
    pub fn new(
        time: Array1<f64>,
        event: Array1<f64>,
        x: Array2<f64>,
        dist: AftDistribution,
    ) -> Result<Self> {
        let n = x.nrows();
        let p = x.ncols();
        if n == 0 || p == 0 {
            return Err(RegressionError::EmptyInput { what: "X" });
        }
        if time.len() != n || event.len() != n {
            return Err(RegressionError::ShapeMismatch {
                what: "time/event length vs X rows",
                expected: n,
                got: time.len().min(event.len()),
            });
        }
        let mut n_events = 0usize;
        for i in 0..n {
            if !time[i].is_finite() || time[i] <= 0.0 {
                return Err(RegressionError::InvalidResponse {
                    msg: format!("survival times must be positive, found {}", time[i]),
                });
            }
            if event[i] == 1.0 {
                n_events += 1;
            } else if event[i] != 0.0 {
                return Err(RegressionError::InvalidResponse {
                    msg: format!("event indicator must be 0 or 1, found {}", event[i]),
                });
            }
        }
        if n_events == 0 {
            return Err(RegressionError::InvalidResponse {
                msg: "no events observed".into(),
            });
        }

        let logt: Vec<f64> = time.iter().map(|t| t.ln()).collect();
        let scale_fixed = dist.scale_fixed();

        // Parameter vector: β (p) then, unless fixed, ln σ.
        let n_par = if scale_fixed { p } else { p + 1 };

        // Negative log-likelihood at parameter vector θ.
        let nll = |theta: &[f64]| -> f64 {
            let sigma = if scale_fixed { 1.0 } else { theta[p].exp() };
            if !sigma.is_finite() || sigma <= 0.0 {
                return f64::INFINITY;
            }
            let mut s = 0.0;
            for i in 0..n {
                let mut eta = 0.0;
                for j in 0..p {
                    eta += x[(i, j)] * theta[j];
                }
                let z = (logt[i] - eta) / sigma;
                if event[i] == 1.0 {
                    s -= -sigma.ln() - logt[i] + dist.log_pdf(z);
                } else {
                    s -= dist.log_surv(z);
                }
            }
            if s.is_finite() {
                s
            } else {
                f64::INFINITY
            }
        };

        // Warm start: OLS of log-time on X (ignoring censoring) for β; residual
        // spread for ln σ.
        let xd = dmatrix_from_rows(n, p, x.as_standard_layout().as_slice().unwrap());
        let ld = dvector_from_slice(&logt);
        let beta0 = match (xd.transpose() * &xd).try_inverse() {
            Some(inv) => inv * xd.transpose() * ld,
            None => return Err(RegressionError::RankDeficient),
        };
        let mut theta0 = vec![0.0; n_par];
        for j in 0..p {
            theta0[j] = beta0[j];
        }
        if !scale_fixed {
            let resid_var = (0..n)
                .map(|i| {
                    let fit: f64 = (0..p).map(|j| x[(i, j)] * beta0[j]).sum();
                    (logt[i] - fit).powi(2)
                })
                .sum::<f64>()
                / n as f64;
            theta0[p] = (0.5 * resid_var.max(1e-6).ln()).max(-5.0);
        }

        let nll_ref = &nll;
        let theta = nelder_mead(nll_ref, &theta0, 0.1, 1e-10, 5000);
        let final_nll = nll(&theta);
        if !final_nll.is_finite() {
            return Err(RegressionError::NotConverged {
                iterations: 5000,
                msg: "AFT optimizer failed to find a finite optimum".into(),
            });
        }

        // Numerical observed information for the covariance.
        let grad = |t: &[f64]| -> Vec<f64> {
            let mut g = vec![0.0; n_par];
            for j in 0..n_par {
                let h = 1e-6 * 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] = (nll(&tp) - nll(&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 hd = dmatrix_from_rows(n_par, n_par, &flat);
        let cov_full = hd.try_inverse().ok_or(RegressionError::RankDeficient)?;

        let coefficients = Array1::from_shape_fn(p, |j| theta[j]);
        let cov = Array2::from_shape_fn((p, p), |(i, j)| cov_full[(i, j)]);
        let scale = if scale_fixed { 1.0 } else { theta[p].exp() };

        Ok(Self {
            dist,
            coefficients,
            scale,
            cov,
            log_likelihood: -final_nll,
            n,
            p,
        })
    }

    /// The assumed distribution.
    pub fn distribution(&self) -> AftDistribution {
        self.dist
    }

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

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

    /// AFT coefficients `β` (log time-acceleration scale), aligned to the design
    /// columns.
    pub fn coefficients(&self) -> ArrayView1<'_, f64> {
        self.coefficients.view()
    }

    /// The scale parameter `σ` (`1` for the exponential).
    pub fn scale(&self) -> f64 {
        self.scale
    }

    /// Coefficient covariance (inverse observed information).
    pub fn covariance(&self) -> ArrayView2<'_, f64> {
        self.cov.view()
    }

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

    /// AIC, `−2ℓ + 2k`, counting the scale parameter unless it is fixed.
    pub fn aic(&self) -> f64 {
        let k = self.p as f64 + if self.dist.scale_fixed() { 0.0 } else { 1.0 };
        -2.0 * self.log_likelihood + 2.0 * k
    }

    /// Coefficient standard errors `√diag(cov)`.
    pub fn coefficient_standard_errors(&self) -> Array1<f64> {
        Array1::from_shape_fn(self.p, |j| self.cov[(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 **median survival time** for a new design matrix `x`:
    /// `exp(xβ + σ·median(W))`.
    pub fn predict_median(&self, x: ArrayView2<'_, f64>) -> Array1<f64> {
        let shift = self.scale * self.dist.median_w();
        x.dot(&self.coefficients).mapv(|eta| (eta + shift).exp())
    }
}