legume-numeric 0.9.3

use super::*;
use crate::matrix::rand_util::collect_f32_seeded;

/// Deterministic LCG for the uniform-draw tests; normal draws come from `normal_chain`.
struct Lcg(u64);
impl Lcg {
    fn next_f32(&mut self) -> f32 {
        self.0 = self
            .0
            .wrapping_mul(6_364_136_223_846_793_005)
            .wrapping_add(1);
        ((self.0 >> 33) as f32) / ((1u64 << 31) as f32) - 0.5
    }
}

#[test]
fn iid_chain_has_ess_near_n() {
    let mut rng = Lcg(42);
    let x: Vec<f32> = (0..4000).map(|_| rng.next_f32()).collect();
    let e = ess(&x);
    // An independent chain is worth ~all of its draws. The estimator is noisy, so this is a
    // band, not an equality — but it must not discount an uncorrelated chain by much.
    assert!(
        e > 0.6 * x.len() as f32,
        "iid chain should keep most of its draws, got ess={e} of {}",
        x.len()
    );
}

#[test]
fn autocorrelated_chain_has_ess_well_below_n() {
    // AR(1) with ρ = 0.9 ⇒ τ = (1+ρ)/(1−ρ) = 19 ⇒ ess ≈ n/19.
    let mut rng = Lcg(7);
    let rho = 0.9f32;
    let mut v = 0.0f32;
    let mut x = Vec::with_capacity(4000);
    for _ in 0..4000 {
        v = rho * v + rng.next_f32();
        x.push(v);
    }
    let e = ess(&x);
    let n = x.len() as f32;
    assert!(
        e < n / 5.0,
        "a strongly autocorrelated chain must be discounted, got ess={e} of {n}"
    );
    assert!(e >= 1.0, "ess stays positive, got {e}");
}

#[test]
fn constant_chain_does_not_divide_by_zero() {
    // The degenerate case the samplers actually produce: every draw on the same side, so the
    // sign-indicator chain is constant. Must return a usable divisor, not NaN/0.
    let e = ess(&vec![1.0f32; 500]);
    assert_eq!(e, 500.0);
    let short = ess(&[1.0, 2.0]);
    assert_eq!(short, 2.0, "too short to estimate ⇒ n");
}

#[test]
fn mcse_is_nonzero_at_p_zero() {
    // The whole point of the Jeffreys smoothing: an lfsr of 0 (no draw on the minority side)
    // must NOT report zero Monte-Carlo error — the plug-in √(p(1−p)/ess) would, and would
    // imply infinite confidence exactly at the top-ranked sites.
    let se = mcse_proportion(0.0, 1000.0);
    assert!(se > 0.0, "p=0 must still carry MC error, got {se}");
    assert!(
        se < 0.002,
        "…but it should be small at ess=1000, got {se}" // ~0.7/ess
    );

    // A borderline lfsr near a 0.1 threshold carries real error at ess=1000: ~0.0095, so the
    // called set genuinely can flip with the seed. This is the number the column exists to show.
    let se_border = mcse_proportion(0.1, 1000.0);
    assert!(
        se_border > 0.008 && se_border < 0.011,
        "lfsr=0.1 at ess=1000 ⇒ mcse≈0.0095, got {se_border}"
    );

    // Fewer effective draws ⇒ strictly more error.
    assert!(mcse_proportion(0.1, 100.0) > se_border);
}

/////////////////
// Split-R̂    //
/////////////////

/// A stationary chain's segments are interchangeable, so R̂ sits at ~1.
#[test]
fn a_stationary_chain_has_rhat_near_one() {
    let x = normal_chain(11, 800, 0.0, 1.0);
    let r = split_rhat(&x);
    assert!(
        (r - 1.0).abs() < 0.05,
        "iid draws should give R̂ ≈ 1, got {r}"
    );
}

/// THE case this exists for: a chain still drifting. The segments then sample different
/// regions, between-segment variance dominates, and R̂ must rise well past the 1.01
/// convention — otherwise a run that never left its initialization would report as
/// converged, which under SGD-xor-sampling is the whole output.
#[test]
fn a_drifting_chain_is_caught() {
    let x: Vec<f32> = (0..800).map(|i| i as f32 * 0.01).collect();
    let r = split_rhat(&x);
    assert!(
        r > 1.1,
        "a linear drift must not pass as converged, got R̂ {r}"
    );
}

/// A chain that never moved is pinned, not divergent — its segments agree exactly. Guards
/// the zero-within-variance branch, which would otherwise divide by zero and report a
/// stuck chain as infinitely bad rather than as suspiciously perfect.
#[test]
fn a_constant_chain_reports_one_not_a_division_by_zero() {
    let x = vec![0.7f32; 400];
    let r = split_rhat(&x);
    assert_eq!(r, 1.0, "a pinned chain should read 1.0, got {r}");
}

/// Segments that are each internally constant but disagree with each other are the one
/// case where zero within-variance means genuine non-convergence, and it must be reported
/// as such rather than silently becoming 1.0.
#[test]
fn constant_but_disagreeing_segments_are_not_converged() {
    let mut x = vec![0.0f32; 200];
    x.extend(vec![5.0f32; 200]);
    assert!(
        split_rhat(&x) > 1.1,
        "two constant halves at different levels are not converged"
    );
}

/// Too short to split ⇒ no claim, rather than a number computed from two points.
#[test]
fn a_chain_too_short_to_split_makes_no_claim() {
    assert_eq!(split_rhat(&[1.0, 2.0, 3.0]), 1.0);
    assert_eq!(split_rhat(&[]), 1.0);
}

/////////////////////////
// split R̂ across chains //
/////////////////////////

fn normal_chain(seed: u64, n: usize, mean: f32, sd: f32) -> Vec<f32> {
    collect_f32_seeded(n, rand_distr::Normal::new(mean, sd).unwrap(), seed)
}

/// Independent chains sampling the same distribution are interchangeable, half by half.
#[test]
fn chains_sampling_the_same_target_have_rhat_near_one() {
    let chains: Vec<Vec<f32>> = (0..4)
        .map(|s| normal_chain(100 + s, 500, 0.0, 1.0))
        .collect();
    let r = split_rhat_chains(&chains);
    assert!(
        (r - 1.0).abs() < 0.05,
        "same target should give R̂ ≈ 1, got {r}"
    );
}

/// THE case this exists for: each chain settled, but in a different mode. Each chain's own
/// [`split_rhat`] is ≈ 1, since both halves agree; only comparing chains sees it.
#[test]
fn chains_stuck_in_different_modes_are_caught() {
    let a = normal_chain(1, 400, 0.0, 0.07);
    let b = normal_chain(2, 400, 2.0, 0.07);
    assert!(split_rhat(&a) < 1.05 && split_rhat(&b) < 1.05);
    let r = split_rhat_chains(&[&a, &b]);
    assert!(r > 2.0, "chains at means 0 and 2 must not pass, got R̂ {r}");
}

/// Chains drifting the same way have equal whole-chain means, so whole-chain
/// Gelman–Rubin calls them converged; the split sees each chain disagree with itself.
#[test]
fn chains_drifting_together_are_caught_by_the_split() {
    let drift = |seed| {
        let mut x = normal_chain(seed, 200, 0.0, 0.1);
        x.extend(normal_chain(seed + 50, 200, 2.0, 0.1));
        x
    };
    let (a, b) = (drift(3), drift(4));
    let whole = gelman_rubin(&[&a, &b]);
    assert!(whole < 1.01, "whole-chain R̂ misses the drift ({whole})");
    let r = split_rhat_chains(&[&a, &b]);
    assert!(r > 1.1, "the split must catch a shared drift, got R̂ {r}");
}

/// Chains of different lengths are compared over the shortest one's draws.
#[test]
fn chains_are_truncated_to_the_shortest() {
    let a = normal_chain(5, 300, 0.0, 1.0);
    let mut b = normal_chain(6, 300, 0.0, 1.0);
    let truncated = split_rhat_chains(&[&a, &b]);
    b.extend(vec![100.0f32; 500]); // past a's length: must be ignored
    assert_eq!(split_rhat_chains(&[&a, &b]), truncated);
}

/// Constant halves follow `split_rhat`: agreeing ⇒ 1, disagreeing ⇒ ∞; too few draws ⇒ 1.
#[test]
fn degenerate_chains_follow_split_rhat() {
    let (a, b) = (vec![0.5f32; 100], vec![0.5f32; 100]);
    assert_eq!(split_rhat_chains(&[&a, &b]), 1.0);
    let c = vec![1.5f32; 100];
    assert_eq!(split_rhat_chains(&[&a, &c]), f32::INFINITY);
    assert_eq!(split_rhat_chains::<&[f32]>(&[]), 1.0);
    assert_eq!(split_rhat_chains(&[&[1.0, 2.0, 3.0][..], &a]), 1.0);
}