legume-numeric 0.9.3

//! Tests for the scalar elliptical slice sampler and its chain runner.

use super::*;
use nalgebra::DVector;
use rand_distr::{Distribution, StandardNormal};

/// Test 1: Gaussian-Gaussian conjugate (1D)
/// Prior: N(0,1), Likelihood: N(y_obs | f, σ²)
/// Posterior: N(μ_post, σ²_post)
#[test]
fn test_conjugate_gaussian_1d() {
    let y_obs = 3.0f32;
    let sigma_sq = 2.0f32; // likelihood variance

    // Analytic posterior
    let sigma_sq_post = 1.0 / (1.0 + 1.0 / sigma_sq);
    let mu_post = sigma_sq_post * y_obs / sigma_sq;

    let lnpdf = move |f: &DVector<f32>| -> f32 {
        let diff = f[0] - y_obs;
        -0.5 * diff * diff / sigma_sq
    };

    let prior_draw = |rng: &mut SmallRng| -> DVector<f32> {
        DVector::from_fn(1, |_, _| {
            let v: f64 = StandardNormal.sample(rng);
            v as f32
        })
    };

    let init = DVector::from_element(1, 0.0f32);

    let sampler = EssSampler {
        n_samples: 10_000,
        warmup: 2_000,
        thin: 1,
        seed: 123,
    };

    let chain = sampler.run(&lnpdf, &prior_draw, &init);
    let mean = chain.posterior_mean();
    let var = chain.posterior_variance();

    assert!(
        (mean[0] - mu_post).abs() < 0.1,
        "mean: {}, expected: {}",
        mean[0],
        mu_post
    );
    assert!(
        (var[0] - sigma_sq_post).abs() < 0.1,
        "var: {}, expected: {}",
        var[0],
        sigma_sq_post
    );
}

/// Test 2: Multivariate Gaussian conjugate (D=5)
/// Prior: N(0, I_5), Likelihood: N(y_obs | f, σ²·I_5)
#[test]
fn test_conjugate_gaussian_5d() {
    let d = 5;
    let sigma_sq = 2.0f32;
    let y_obs = DVector::from_vec(vec![1.0, -2.0, 0.5, 3.0, -1.5]);

    let sigma_sq_post = 1.0 / (1.0 + 1.0 / sigma_sq);

    let lnpdf = {
        let y = y_obs.clone();
        move |f: &DVector<f32>| -> f32 {
            let diff = f - &y;
            -0.5 * diff.dot(&diff) / sigma_sq
        }
    };

    let prior_draw = move |rng: &mut SmallRng| -> DVector<f32> {
        DVector::from_fn(d, |_, _| {
            let v: f64 = StandardNormal.sample(rng);
            v as f32
        })
    };

    let init = DVector::from_element(d, 0.0f32);

    let sampler = EssSampler {
        n_samples: 10_000,
        warmup: 2_000,
        thin: 1,
        seed: 456,
    };

    let chain = sampler.run(&lnpdf, &prior_draw, &init);
    let mean = chain.posterior_mean();

    for j in 0..d {
        let analytic = sigma_sq_post * y_obs[j] / sigma_sq;
        assert!(
            (mean[j] - analytic).abs() < 0.15,
            "dim {}: mean={}, expected={}",
            j,
            mean[j],
            analytic
        );
    }
}

/// Test 3: Correlated prior (D=3) with Cholesky
/// Prior: N(0, Σ), Likelihood: N(y_obs | f, σ²·I)
/// Posterior: N(Σ_post · y/σ², Σ_post) where Σ_post = (Σ⁻¹ + I/σ²)⁻¹
#[test]
fn test_correlated_prior_3d() {
    use nalgebra::DMatrix;

    let d = 3;
    let sigma_sq = 1.5f32;
    let y_obs = DVector::from_vec(vec![2.0, -1.0, 1.5]);

    // Prior covariance with off-diagonal correlations
    #[rustfmt::skip]
    let sigma_prior = DMatrix::from_row_slice(d, d, &[
        1.0, 0.5, 0.2,
        0.5, 1.0, 0.3,
        0.2, 0.3, 1.0,
    ]);

    // Cholesky decomposition of prior covariance
    let chol = sigma_prior.clone().cholesky().unwrap();
    let chol_l = chol.l();

    // Analytic posterior: Σ_post = (Σ⁻¹ + I/σ²)⁻¹
    let sigma_prior_inv = sigma_prior.clone().try_inverse().unwrap();
    let eye = DMatrix::identity(d, d);
    let precision_post = &sigma_prior_inv + &eye / sigma_sq;
    let sigma_post = precision_post.try_inverse().unwrap();
    let mu_post = &sigma_post * &y_obs / sigma_sq;

    let lnpdf = {
        let y = y_obs.clone();
        move |f: &DVector<f32>| -> f32 {
            let diff = f - &y;
            -0.5 * diff.dot(&diff) / sigma_sq
        }
    };

    let prior_draw = {
        let l = chol_l.clone();
        move |rng: &mut SmallRng| -> DVector<f32> {
            let z = DVector::from_fn(d, |_, _| {
                let v: f64 = StandardNormal.sample(rng);
                v as f32
            });
            &l * z
        }
    };

    let init = DVector::from_element(d, 0.0f32);

    let sampler = EssSampler {
        n_samples: 20_000,
        warmup: 5_000,
        thin: 1,
        seed: 789,
    };

    let chain = sampler.run(&lnpdf, &prior_draw, &init);
    let mean = chain.posterior_mean();

    for j in 0..d {
        assert!(
            (mean[j] - mu_post[j]).abs() < 0.15,
            "dim {}: mean={}, expected={}",
            j,
            mean[j],
            mu_post[j]
        );
    }
}

/// Test 4: Non-conjugate logistic regression
/// Prior: N(0, I), Likelihood: logistic
/// Check posterior mean is in the right region (near MAP).
#[test]
fn test_logistic_regression() {
    let d = 2;

    // Simple dataset: 4 points
    let x = vec![
        DVector::from_vec(vec![1.0, 0.5]),
        DVector::from_vec(vec![-1.0, -0.3]),
        DVector::from_vec(vec![0.8, -0.7]),
        DVector::from_vec(vec![-0.5, 1.0]),
    ];
    let y = vec![1.0f32, 0.0, 1.0, 0.0]; // labels

    let lnpdf = {
        let x = x.clone();
        let y = y.clone();
        move |f: &DVector<f32>| -> f32 {
            let mut ll = 0.0f32;
            for (xi, &yi) in x.iter().zip(y.iter()) {
                let eta = f.dot(xi);
                // log p(y|eta) = y*eta - log(1+exp(eta))
                let log1pexp = if eta > 20.0 {
                    eta
                } else if eta < -20.0 {
                    0.0
                } else {
                    (1.0 + eta.exp()).ln()
                };
                ll += yi * eta - log1pexp;
            }
            ll
        }
    };

    let prior_draw = move |rng: &mut SmallRng| -> DVector<f32> {
        DVector::from_fn(d, |_, _| {
            let v: f64 = StandardNormal.sample(rng);
            v as f32
        })
    };

    let init = DVector::from_element(d, 0.0f32);

    let sampler = EssSampler {
        n_samples: 10_000,
        warmup: 2_000,
        thin: 1,
        seed: 101,
    };

    let chain = sampler.run(&lnpdf, &prior_draw, &init);
    let mean = chain.posterior_mean();

    // With this data and N(0,I) prior, the posterior mean should be
    // in the positive region for dim 0 (positive examples have positive x[0]).
    // Just verify we get a reasonable, non-degenerate result.
    assert!(
        mean[0] > 0.0,
        "expected positive posterior mean for dim 0, got {}",
        mean[0]
    );
    assert!(
        mean[0].abs() < 5.0 && mean[1].abs() < 5.0,
        "posterior mean should be moderate: {:?}",
        mean
    );
}

/// Test 5: Parallel chains give consistent results
#[test]
fn test_parallel_chains() {
    let y_obs = 2.0f32;
    let sigma_sq = 1.0f32;
    let sigma_sq_post = 1.0 / (1.0 + 1.0 / sigma_sq);
    let mu_post = sigma_sq_post * y_obs / sigma_sq;

    let lnpdf = move |f: &DVector<f32>| -> f32 {
        let diff = f[0] - y_obs;
        -0.5 * diff * diff / sigma_sq
    };

    let prior_draw = |rng: &mut SmallRng| -> DVector<f32> {
        DVector::from_fn(1, |_, _| {
            let v: f64 = StandardNormal.sample(rng);
            v as f32
        })
    };

    let init = DVector::from_element(1, 0.0f32);

    let sampler = EssSampler {
        n_samples: 5_000,
        warmup: 1_000,
        thin: 1,
        seed: 222,
    };

    let chains = sampler.run_parallel(4, &lnpdf, &prior_draw, &init);
    assert_eq!(chains.len(), 4);

    // Each chain's posterior mean should be close to analytic
    for (i, chain) in chains.iter().enumerate() {
        let mean = chain.posterior_mean();
        assert!(
            (mean[0] - mu_post).abs() < 0.2,
            "chain {}: mean={}, expected={}",
            i,
            mean[0],
            mu_post
        );
    }
}

/// Test 6: Candle Tensor — Gaussian conjugate (D=5)
/// Same setup as test 2, but using Tensor on CPU.
#[test]
#[cfg(feature = "tensor")]
fn test_tensor_conjugate_gaussian_5d() {
    use candle_core::{DType, Device, Tensor};

    let device = Device::Cpu;
    let d = 5;
    let sigma_sq = 2.0f32;
    let y_vals = vec![1.0f32, -2.0, 0.5, 3.0, -1.5];

    let sigma_sq_post = 1.0 / (1.0 + 1.0 / sigma_sq);

    let y_tensor = Tensor::from_vec(y_vals.clone(), &[d], &device).unwrap();

    let lnpdf = {
        let y = y_tensor.clone();
        move |f: &Tensor| -> f32 {
            let diff = (f - &y).unwrap();
            let sq = (&diff * &diff).unwrap().sum_all().unwrap();
            (sq.to_scalar::<f32>().unwrap()) * (-0.5 / sigma_sq)
        }
    };

    let prior_draw = {
        let dev = device.clone();
        move |_rng: &mut SmallRng| -> Tensor { Tensor::randn(0f32, 1f32, &[d], &dev).unwrap() }
    };

    let init = Tensor::zeros(&[d], DType::F32, &device).unwrap();

    let sampler = EssSampler {
        n_samples: 10_000,
        warmup: 2_000,
        thin: 1,
        seed: 456,
    };

    let chain = sampler.run(&lnpdf, &prior_draw, &init);

    // Manual posterior mean (no EssParamSummary for Tensor)
    let n = chain.n_samples();
    let stacked = Tensor::stack(&chain.samples, 0).unwrap();
    let mean_tensor = (stacked.sum(0).unwrap() / n as f64).unwrap();
    let mean: Vec<f32> = mean_tensor.to_vec1().unwrap();

    for j in 0..d {
        let analytic = sigma_sq_post * y_vals[j] / sigma_sq;
        assert!(
            (mean[j] - analytic).abs() < 0.15,
            "dim {}: mean={}, expected={}",
            j,
            mean[j],
            analytic
        );
    }
}

/// Test 7: Composite parameters via Vec<DVector<f32>>
/// Two independent blocks: β (D=3) and γ (D=2) with separate likelihoods.
/// Prior: N(0, I) on each block independently.
/// Composite log-likelihood = llik_1(β) + llik_2(γ).
#[test]
fn test_composite_vec_params() {
    let y1 = DVector::from_vec(vec![1.0f32, -0.5, 2.0]); // observed for block 1
    let y2 = DVector::from_vec(vec![-1.0f32, 3.0]); // observed for block 2
    let sigma_sq_1 = 1.0f32;
    let sigma_sq_2 = 2.0f32;

    // Analytic posteriors (independent conjugate Gaussians)
    let sp1 = 1.0 / (1.0 + 1.0 / sigma_sq_1);
    let sp2 = 1.0 / (1.0 + 1.0 / sigma_sq_2);

    let lnpdf = {
        let y1 = y1.clone();
        let y2 = y2.clone();
        move |params: &Vec<DVector<f32>>| -> f32 {
            let diff1 = &params[0] - &y1;
            let diff2 = &params[1] - &y2;
            -0.5 * diff1.dot(&diff1) / sigma_sq_1 + -0.5 * diff2.dot(&diff2) / sigma_sq_2
        }
    };

    let prior_draw = |rng: &mut SmallRng| -> Vec<DVector<f32>> {
        vec![
            DVector::from_fn(3, |_, _| {
                let v: f64 = StandardNormal.sample(rng);
                v as f32
            }),
            DVector::from_fn(2, |_, _| {
                let v: f64 = StandardNormal.sample(rng);
                v as f32
            }),
        ]
    };

    let init = vec![
        DVector::from_element(3, 0.0f32),
        DVector::from_element(2, 0.0f32),
    ];

    let sampler = EssSampler {
        n_samples: 10_000,
        warmup: 2_000,
        thin: 1,
        seed: 333,
    };

    let chain = sampler.run(&lnpdf, &prior_draw, &init);
    let n = chain.n_samples() as f32;

    // Compute mean for each block manually
    let mut mean1 = DVector::from_element(3, 0.0f32);
    let mut mean2 = DVector::from_element(2, 0.0f32);
    for sample in &chain.samples {
        mean1 += &sample[0];
        mean2 += &sample[1];
    }
    mean1 /= n;
    mean2 /= n;

    for j in 0..3 {
        let analytic = sp1 * y1[j] / sigma_sq_1;
        assert!(
            (mean1[j] - analytic).abs() < 0.15,
            "block1 dim {}: mean={}, expected={}",
            j,
            mean1[j],
            analytic
        );
    }
    for j in 0..2 {
        let analytic = sp2 * y2[j] / sigma_sq_2;
        assert!(
            (mean2[j] - analytic).abs() < 0.15,
            "block2 dim {}: mean={}, expected={}",
            j,
            mean2[j],
            analytic
        );
    }
}

/// A scalar `f32` parameter is the 1-element `DVector` without the allocation: same RNG
/// draws, same arithmetic, so the chains match draw for draw.
#[test]
fn scalar_param_matches_one_element_vector() {
    let y_obs = 3.0f32;
    let sigma_sq = 2.0f32;
    let normal = |rng: &mut SmallRng| -> f32 {
        let v: f64 = StandardNormal.sample(rng);
        v as f32
    };

    let sampler = EssSampler {
        n_samples: 2_000,
        warmup: 500,
        thin: 2,
        seed: 7,
    };
    let scalar = sampler.run(
        &|f: &f32| -0.5 * (f - y_obs) * (f - y_obs) / sigma_sq,
        &normal,
        &0.0f32,
    );
    let vector = sampler.run(
        &|f: &DVector<f32>| -0.5 * (f[0] - y_obs) * (f[0] - y_obs) / sigma_sq,
        &|rng: &mut SmallRng| DVector::from_element(1, normal(rng)),
        &DVector::from_element(1, 0.0f32),
    );

    let vector_draws: Vec<f32> = vector.samples.iter().map(|v| v[0]).collect();
    assert_eq!(scalar.samples, vector_draws);
    assert_eq!(scalar.posterior_mean(), vector.posterior_mean());
}