regression-diagnostics 0.2.0

use ndarray::Array1;

use crate::error::{RegressionError, Result};

/// One step of a **Kaplan–Meier** survival curve, recorded at a distinct event
/// time.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct KmStep {
    /// The event time.
    pub time: f64,
    /// Number at risk just before this time, `nᵢ`.
    pub at_risk: usize,
    /// Number of events at this time, `dᵢ`.
    pub events: usize,
    /// Survivor function `Ŝ(t)` at and after this time.
    pub survival: f64,
    /// Greenwood standard error of `Ŝ(t)`.
    pub std_error: f64,
}

/// The **Kaplan–Meier** product-limit estimator of the survivor function
/// `S(t) = P(T > t)` from right-censored data.
///
/// `Ŝ(t) = Π_{tᵢ ≤ t} (1 − dᵢ/nᵢ)` over distinct event times, with `dᵢ` events
/// and `nᵢ` subjects at risk. Censored observations leave the risk set without
/// producing a step. Standard errors use **Greenwood's formula**. This is the
/// non-parametric, covariate-free counterpart to the Cox model's baseline.
#[derive(Debug, Clone)]
pub struct KaplanMeier {
    steps: Vec<KmStep>,
    n: usize,
}

impl KaplanMeier {
    /// Estimate the survival curve from event/censoring `time` and `event`
    /// indicators (`1.0` event, `0.0` right-censored).
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] if there are no observations.
    /// * [`RegressionError::ShapeMismatch`] if `time` and `event` differ in
    ///   length.
    /// * [`RegressionError::InvalidResponse`] for non-positive times or event
    ///   flags outside `{0, 1}`.
    pub fn new(time: Array1<f64>, event: Array1<f64>) -> Result<Self> {
        let n = time.len();
        if n == 0 {
            return Err(RegressionError::EmptyInput { what: "time" });
        }
        if event.len() != n {
            return Err(RegressionError::ShapeMismatch {
                what: "event length vs time",
                expected: n,
                got: event.len(),
            });
        }
        for &t in time.iter() {
            if !t.is_finite() || t <= 0.0 {
                return Err(RegressionError::InvalidResponse {
                    msg: format!("survival times must be positive, found {t}"),
                });
            }
        }
        for &e in event.iter() {
            if e != 0.0 && e != 1.0 {
                return Err(RegressionError::InvalidResponse {
                    msg: format!("event indicator must be 0 or 1, found {e}"),
                });
            }
        }

        // Distinct event times, ascending.
        let mut ev_times: Vec<f64> = time
            .iter()
            .zip(event.iter())
            .filter(|&(_, &e)| e == 1.0)
            .map(|(&t, _)| t)
            .collect();
        ev_times.sort_by(|a, b| a.partial_cmp(b).unwrap());
        ev_times.dedup();

        let mut steps = Vec::with_capacity(ev_times.len());
        let mut survival = 1.0;
        let mut greenwood_sum = 0.0; // Σ dᵢ / (nᵢ (nᵢ − dᵢ))
        for &t in &ev_times {
            let at_risk = time.iter().filter(|&&tj| tj >= t).count();
            let events = time
                .iter()
                .zip(event.iter())
                .filter(|&(&tj, &e)| tj == t && e == 1.0)
                .count();
            let ni = at_risk as f64;
            let di = events as f64;
            survival *= 1.0 - di / ni;
            if ni > di {
                greenwood_sum += di / (ni * (ni - di));
            }
            let std_error = survival * greenwood_sum.sqrt();
            steps.push(KmStep {
                time: t,
                at_risk,
                events,
                survival,
                std_error,
            });
        }

        Ok(Self { steps, n })
    }

    /// Number of observations the estimator was built from.
    pub fn n_observations(&self) -> usize {
        self.n
    }

    /// The estimated steps, one per distinct event time (ascending).
    pub fn steps(&self) -> &[KmStep] {
        &self.steps
    }

    /// Survivor function `Ŝ(t)` at an arbitrary time (right-continuous step
    /// function; `1.0` before the first event).
    pub fn survival_at(&self, t: f64) -> f64 {
        let mut s = 1.0;
        for step in &self.steps {
            if step.time <= t {
                s = step.survival;
            } else {
                break;
            }
        }
        s
    }

    /// Median survival time — the earliest event time at which `Ŝ(t) ≤ 0.5`, or
    /// `None` if the curve never drops to `0.5` (heavy censoring).
    pub fn median_survival(&self) -> Option<f64> {
        self.steps
            .iter()
            .find(|s| s.survival <= 0.5)
            .map(|s| s.time)
    }
}