regression-diagnostics 0.2.0

//! Survival diagnostics. The strongest anchor is external: the Cox fit on the
//! classic Freireich (1963) leukemia remission data reproduces the published
//! `coxph` coefficients for both tie methods — Breslow ≈ 1.509 (se ≈ 0.410) and
//! Efron ≈ 1.572 (se ≈ 0.412), bracketed by the exact estimate ≈ 1.628 (Efron
//! corrects Breslow's bias toward zero, so its magnitude is the larger).
//! Structural anchors pin down the rest — the score vanishes at the MLE (so
//! martingale residuals sum to zero, and Breslow Schoenfeld residuals sum to
//! zero at the Breslow fit), concordance lies in [0, 1], and the Kaplan–Meier
//! steps match the product-limit values computed by hand.

use ndarray::{Array1, Array2};
use regression_diagnostics::survival::{
    deviance_residuals, martingale_residuals, schoenfeld_residuals, AftDistribution, AftFit, CoxFit,
    KaplanMeier, Ties,
};
use regression_diagnostics::RegressionError;

/// Freireich 6-MP vs placebo remission times (weeks). Covariate x = 1 for the
/// placebo group (higher hazard). Returns (time, event, X).
fn freireich() -> (Array1<f64>, Array1<f64>, Array2<f64>) {
    // Placebo group: all 21 are events.
    let placebo_t = [
        1.0, 1.0, 2.0, 2.0, 3.0, 4.0, 4.0, 5.0, 5.0, 8.0, 8.0, 8.0, 8.0, 11.0, 11.0, 12.0, 12.0,
        15.0, 17.0, 22.0, 23.0,
    ];
    // 6-MP group: 9 events then 12 right-censored.
    let mp_t = [
        6.0, 6.0, 6.0, 7.0, 10.0, 13.0, 16.0, 22.0, 23.0, // events
        6.0, 9.0, 10.0, 11.0, 17.0, 19.0, 20.0, 25.0, 32.0, 32.0, 34.0, 35.0, // censored
    ];
    let mp_e = [
        1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
        0.0, 0.0, 0.0,
    ];

    let n = placebo_t.len() + mp_t.len();
    let mut time = Array1::<f64>::zeros(n);
    let mut event = Array1::<f64>::zeros(n);
    let mut x = Array2::<f64>::zeros((n, 1));
    let mut r = 0;
    for &t in placebo_t.iter() {
        time[r] = t;
        event[r] = 1.0;
        x[(r, 0)] = 1.0; // placebo
        r += 1;
    }
    for (k, &t) in mp_t.iter().enumerate() {
        time[r] = t;
        event[r] = mp_e[k];
        x[(r, 0)] = 0.0; // 6-MP
        r += 1;
    }
    (time, event, x)
}

#[test]
fn cox_matches_freireich_reference() {
    let (time, event, x) = freireich();

    // Default (Efron) estimate: β ≈ 1.5721, SE ≈ 0.4124, HR ≈ 4.817.
    let efron = CoxFit::new(time.clone(), event.clone(), x.clone()).unwrap();
    assert!(
        (efron.coefficients()[0] - 1.5721).abs() < 0.01,
        "Efron β = {}, expected ≈ 1.5721",
        efron.coefficients()[0]
    );
    assert!((efron.coefficient_standard_errors()[0] - 0.4124).abs() < 0.01);
    assert!((efron.hazard_ratios()[0] - 4.817).abs() < 0.05);
    assert!(efron.p_values()[0] < 1e-3);

    // Breslow estimate: β ≈ 1.5092, SE ≈ 0.4096 — the smaller magnitude.
    let breslow =
        CoxFit::with_options(time, event, x, Ties::Breslow, 100, 1e-9).unwrap();
    assert!(
        (breslow.coefficients()[0] - 1.5092).abs() < 0.01,
        "Breslow β = {}, expected ≈ 1.5092",
        breslow.coefficients()[0]
    );
    assert!((breslow.coefficient_standard_errors()[0] - 0.4096).abs() < 0.01);
    assert!(breslow.coefficients()[0] < efron.coefficients()[0]);
}

#[test]
fn cox_score_zero_shows_as_residual_sums() {
    let (time, event, x) = freireich();
    // Martingale residuals sum to zero for any fit (the Breslow baseline is
    // calibrated so Σ Ĥᵢ equals the event count) — check on the Efron default.
    let efron = CoxFit::new(time.clone(), event.clone(), x.clone()).unwrap();
    let m = martingale_residuals(&efron);
    assert!(m.sum().abs() < 1e-6, "Σ martingale = {}", m.sum());
    assert!(deviance_residuals(&efron).iter().all(|v| v.is_finite()));
    let c = efron.concordance();
    assert!((0.0..=1.0).contains(&c) && c > 0.5, "C-index = {c}");

    // The classic Schoenfeld residual uses the full-risk-set mean, so its
    // per-covariate sum is the Breslow score — zero exactly at the Breslow fit.
    let breslow =
        CoxFit::with_options(time, event, x, Ties::Breslow, 100, 1e-9).unwrap();
    let (_times, sch) = schoenfeld_residuals(&breslow);
    for a in 0..breslow.n_parameters() {
        let col_sum: f64 = (0..sch.nrows()).map(|r| sch[(r, a)]).sum();
        assert!(col_sum.abs() < 1e-6, "Σ schoenfeld[{a}] = {col_sum}");
    }
}

#[test]
fn efron_and_breslow_agree_without_ties() {
    // Distinct times ⇒ Efron and Breslow coincide exactly.
    let n = 12usize;
    let mut time = Array1::<f64>::zeros(n);
    let mut event = Array1::<f64>::zeros(n);
    let mut x = Array2::<f64>::zeros((n, 1));
    for i in 0..n {
        time[i] = (i + 1) as f64; // all distinct
        event[i] = if i % 4 == 0 { 0.0 } else { 1.0 };
        x[(i, 0)] = (i as f64) * 0.3 - 1.5;
    }
    let efron = CoxFit::with_options(time.clone(), event.clone(), x.clone(), Ties::Efron, 100, 1e-9)
        .unwrap();
    let breslow = CoxFit::with_options(time, event, x, Ties::Breslow, 100, 1e-9).unwrap();
    assert!((efron.coefficients()[0] - breslow.coefficients()[0]).abs() < 1e-9);
}

#[test]
fn kaplan_meier_matches_hand_computation() {
    // time = [2,3,4,4,5,6], event = [1,1,1,0,1,0].
    // Ŝ at event times 2,3,4,5 = 5/6, 2/3, 1/2, 1/4.
    let time = Array1::from(vec![2.0, 3.0, 4.0, 4.0, 5.0, 6.0]);
    let event = Array1::from(vec![1.0, 1.0, 1.0, 0.0, 1.0, 0.0]);
    let km = KaplanMeier::new(time, event).unwrap();
    let steps = km.steps();
    assert_eq!(steps.len(), 4);
    let want = [5.0 / 6.0, 2.0 / 3.0, 1.0 / 2.0, 1.0 / 4.0];
    let want_at_risk = [6, 5, 4, 2];
    for (k, step) in steps.iter().enumerate() {
        assert!((step.survival - want[k]).abs() < 1e-12, "S step {k}");
        assert_eq!(step.at_risk, want_at_risk[k]);
        assert!(step.std_error >= 0.0);
    }
    // survival_at is a right-continuous step function.
    assert_eq!(km.survival_at(1.0), 1.0);
    assert!((km.survival_at(3.5) - 2.0 / 3.0).abs() < 1e-12);
    assert_eq!(km.median_survival(), Some(4.0)); // first time S <= 0.5
}

#[test]
fn cox_baseline_hazard_is_increasing() {
    let (time, event, x) = freireich();
    let fit = CoxFit::new(time, event, x).unwrap();
    let h0 = fit.baseline_cumulative_hazard();
    assert!(!h0.is_empty());
    for w in h0.windows(2) {
        assert!(w[1].1 >= w[0].1, "cumulative hazard must be non-decreasing");
        assert!(w[1].0 > w[0].0, "times must be strictly increasing");
    }
}

#[test]
fn exponential_aft_intercept_matches_closed_form() {
    // For an intercept-only exponential model the MLE rate is (#events)/(Σ times),
    // so the AFT intercept β₀ = ln(Σt / d).
    let time = Array1::from(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]);
    let event = Array1::from(vec![1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0]); // 6 events
    let x = Array2::<f64>::ones((8, 1));
    let fit = AftFit::new(time, event, x, AftDistribution::Exponential).unwrap();
    let want = (36.0_f64 / 6.0).ln(); // Σt = 36, d = 6
    assert!(
        (fit.coefficients()[0] - want).abs() < 1e-3,
        "β₀ = {}, expected ln 6 = {want}",
        fit.coefficients()[0]
    );
    assert_eq!(fit.scale(), 1.0); // exponential fixes σ
}

#[test]
fn weibull_aft_fits_and_predicts() {
    // A covariate that accelerates failure time; Weibull estimates a scale.
    let n = 30usize;
    let mut time = Array1::<f64>::zeros(n);
    let mut event = Array1::<f64>::zeros(n);
    let mut x = Array2::<f64>::ones((n, 2));
    for i in 0..n {
        let xi = (i % 2) as f64; // two groups
        x[(i, 1)] = xi;
        // group 1 fails later (positive AFT coefficient)
        time[i] = 2.0 + xi * 3.0 + (i as f64 % 5.0) * 0.5;
        event[i] = if i % 6 == 0 { 0.0 } else { 1.0 };
    }
    let fit = AftFit::new(time, event, x.clone(), AftDistribution::Weibull).unwrap();
    assert!(fit.log_likelihood().is_finite());
    assert!(fit.scale() > 0.0);
    assert!(fit.coefficient_standard_errors().iter().all(|v| v.is_finite()));
    // Group 1's predicted median survival exceeds group 0's.
    let med = fit.predict_median(x.view());
    assert!(med.iter().all(|m| *m > 0.0));
    assert!(fit.coefficients()[1] > 0.0, "accelerant coef should be positive");
}

#[test]
fn stratified_with_one_stratum_equals_ordinary() {
    let n = 16usize;
    let mut time = Array1::<f64>::zeros(n);
    let mut event = Array1::<f64>::zeros(n);
    let mut x = Array2::<f64>::zeros((n, 1));
    for i in 0..n {
        time[i] = (i + 1) as f64;
        event[i] = if i % 5 == 0 { 0.0 } else { 1.0 };
        x[(i, 0)] = (i as f64) * 0.2 - 1.5;
    }
    let ordinary = CoxFit::new(time.clone(), event.clone(), x.clone()).unwrap();
    let one_stratum = vec![7usize; n]; // all in the same (arbitrary) stratum
    let strat = CoxFit::stratified(time, event, x, &one_stratum, Ties::Efron).unwrap();
    assert!((ordinary.coefficients()[0] - strat.coefficients()[0]).abs() < 1e-9);
    assert_eq!(strat.n_strata(), 1);
}

#[test]
fn counting_process_single_intervals_equals_ordinary() {
    // (0, tᵢ] intervals reproduce the ordinary right-censored Cox fit.
    let n = 16usize;
    let mut stop = Array1::<f64>::zeros(n);
    let mut start = Array1::<f64>::zeros(n);
    let mut event = Array1::<f64>::zeros(n);
    let mut x = Array2::<f64>::zeros((n, 1));
    for i in 0..n {
        stop[i] = (i + 1) as f64;
        start[i] = 0.0;
        event[i] = if i % 5 == 0 { 0.0 } else { 1.0 };
        x[(i, 0)] = (i as f64) * 0.2 - 1.5;
    }
    let ordinary = CoxFit::new(stop.clone(), event.clone(), x.clone()).unwrap();
    let cp = CoxFit::counting_process(start, stop, event, x, Ties::Efron).unwrap();
    assert!((ordinary.coefficients()[0] - cp.coefficients()[0]).abs() < 1e-9);
}

#[test]
fn counting_process_rejects_bad_intervals() {
    let start = Array1::from(vec![0.0, 2.0, 1.0]);
    let stop = Array1::from(vec![1.0, 2.0, 3.0]); // interval 1 has start == stop
    let event = Array1::from(vec![1.0, 1.0, 0.0]);
    let mut x = Array2::<f64>::zeros((3, 1));
    for i in 0..3 {
        x[(i, 0)] = i as f64;
    }
    assert!(matches!(
        CoxFit::counting_process(start, stop, event, x, Ties::Efron).unwrap_err(),
        RegressionError::InvalidResponse { .. }
    ));
}

#[test]
fn cox_rejects_bad_input() {
    let time = Array1::from(vec![1.0, 2.0, 3.0, 4.0]);
    // Intercept column present.
    let x_const = Array2::<f64>::ones((4, 1));
    let ev = Array1::from(vec![1.0, 0.0, 1.0, 1.0]);
    assert!(matches!(
        CoxFit::new(time.clone(), ev.clone(), x_const).unwrap_err(),
        RegressionError::InvalidResponse { .. }
    ));
    // No events at all.
    let mut x = Array2::<f64>::zeros((4, 1));
    for i in 0..4 {
        x[(i, 0)] = i as f64;
    }
    let no_events = Array1::from(vec![0.0, 0.0, 0.0, 0.0]);
    assert!(matches!(
        CoxFit::new(time, no_events, x).unwrap_err(),
        RegressionError::InvalidResponse { .. }
    ));
}