prodef 0.2.2

//! Comprehensive testing module for ProDeF distributions.
//!
//! All tests for probability distributions are contained in this module.
//! Tests follow a consistent pattern: constructor validation, sampling behavior,
//! density evaluation, and domain constraint checking.
//!
//! # Test Organization
//!
//! - **Constructor Tests**: Validate parameter acceptance/rejection
//! - **Sampling Tests**: Check determinism, convergence, and domain containment
//! - **Density Tests**: Verify PDF evaluation and None handling
//! - **Domain Tests**: Confirm boundary behavior and containment
//!
//! # Distributions Tested
//!
//! **Univariate:**
//! - Constant: Fixed value
//! - Uniform: Bounded uniform
//! - Normal: Gaussian (unbounded and bounded)
//! - Cosine: Cosine-shaped
//! - Lognormal: Log-normal
//! - LogUniform: Log-uniform
//!
//! **Multivariate:**
//! - MultivariateDensity: Product of independent univariates
//! - MultivariateNormalDensity: Gaussian with full covariance matrix
//! - ParticleDensity: Kernel density estimation from weighted particles

use crate::{
    ConstantDensity, CosineDensity, Density, Domain, LogUniformDensity, LognormalDensity,
    MultivariateDensity, MultivariateNormalDensity, NormalDensity, ParticleDensity, SamplingMode,
    UniformDensity, macros::tval,
};
use approx::assert_abs_diff_eq;
use nalgebra::{
    DefaultAllocator, Dyn, Matrix2, OVector, RealField, SVector, U1, U2, allocator::Allocator,
};
use rand::{SeedableRng, rngs::Xoshiro256PlusPlus};

// ============================================================================
// CONFIGURATION CONSTANTS
// ============================================================================

/// Large sample size for statistical convergence tests
const N_SAMPLES: usize = 10_000;

/// Default RNG seed for reproducible tests
const RNG_SEED: u64 = 42;

// ============================================================================
// HELPER FUNCTIONS
// ============================================================================

/// Asserts density evaluation within and outside domain.
fn assert_density_correct<'a, D>(dist: &'a D, sample: f64, should_exist: bool)
where
    &'a D: Density<f64, U1>,
{
    let sample_vec = SVector::from([sample]);
    let density = dist.density::<U1, U1>(&sample_vec.as_view());

    if should_exist {
        assert!(density.is_some(), "Expected density to exist at {}", sample);
        assert!(density.unwrap() > 0.0, "Density should be positive");
    } else {
        assert!(
            density.is_none(),
            "Expected density to be None at {}",
            sample
        );
    }
}

/// Asserts all samples fall within bounds.
fn assert_samples_in_bounds<T: PartialOrd + std::fmt::Display>(
    samples: &[T],
    min: T,
    max: T,
    inclusive: bool,
) {
    for sample in samples {
        let in_bounds = if inclusive {
            sample >= &min && sample <= &max
        } else {
            sample > &min && sample < &max
        };
        assert!(
            in_bounds,
            "Sample {} outside bounds [{}, {}]",
            sample, min, max
        );
    }
}

/// Asserts sampling determinism: same seed produces identical samples.
fn assert_sampling_determinism<'a, D, T, Dim>(dist: &'a D, seed: u64, n_samples: usize)
where
    T: RealField + PartialEq + std::fmt::Debug,
    Dim: nalgebra::Dim,
    DefaultAllocator: Allocator<Dim>,
    &'a D: Density<T, Dim>,
{
    let samples1 = collect_samples_nd(dist, n_samples, seed);
    let samples2 = collect_samples_nd(dist, n_samples, seed);

    assert_eq!(
        samples1, samples2,
        "Sampling not deterministic with same seed"
    );
}

/// Collects N samples from an N-dimensional distribution with deterministic seeding.
/// Returns Vec<OVector<T, D>> where OVector contains the D-dimensional samples.
/// This is the canonical collection function; use it for all distributions.
fn collect_samples_nd<'a, T, D, G>(dist: &'a G, n_samples: usize, seed: u64) -> Vec<OVector<T, D>>
where
    T: RealField,
    D: nalgebra::Dim,
    &'a G: Density<T, D>,
    DefaultAllocator: Allocator<D>,
{
    let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed);
    let mode = SamplingMode::default();

    (0..n_samples)
        .filter_map(|_| dist.sample(&mut rng, &mode))
        .collect()
}

/// Convenience wrapper: Collects N samples from a 1D distribution and extracts scalar values.
/// Uses collect_samples_nd internally. Returns Vec<T> where T is the scalar type.
fn collect_samples<'a, T, D>(dist: &'a D, n_samples: usize, seed: u64) -> Vec<T>
where
    T: RealField,
    &'a D: Density<T, U1>,
{
    let samples = collect_samples_nd(dist, n_samples, seed);
    samples.into_iter().map(|s| s[0].clone()).collect()
}

/// Computes the mean of a vector of samples.
fn compute_sample_mean<T: RealField>(samples: &[T]) -> T {
    if samples.is_empty() {
        return T::zero();
    }

    let sum = samples.iter().cloned().fold(T::zero(), T::add);
    sum / tval!(samples.len(), usize)
}

/// Computes the sample variance from samples and mean.
fn compute_sample_variance<T: RealField>(samples: &[T], mean: T) -> T {
    if samples.len() < 2 {
        return T::zero();
    }

    let sum_sq_diff = samples
        .iter()
        .cloned()
        .fold(T::zero(), |acc, x| acc + (x - mean.clone()).powi(2));

    sum_sq_diff / tval!(samples.len() - 1, usize)
}

/// Computes the mean vector of N-D samples.
/// For D-dimensional samples, returns the arithmetic mean across all samples.
#[allow(dead_code)]
fn compute_sample_mean_nd<T: RealField, D: nalgebra::Dim>(
    samples: &[OVector<T, D>],
) -> OVector<T, D>
where
    DefaultAllocator: Allocator<D>,
{
    if samples.is_empty() {
        return OVector::zeros_generic(samples[0].shape_generic().0, U1);
    }

    let sum = samples.iter().cloned().fold(
        OVector::zeros_generic(samples[0].shape_generic().0, U1),
        |acc, x| acc + x,
    );
    sum / tval!(samples.len(), usize)
}

/// Extracts a single dimension from N-D samples.
/// Returns a vector of scalar values for the specified dimension index.
#[allow(dead_code)]
fn extract_dimension<T: RealField, D: nalgebra::Dim>(
    samples: &[OVector<T, D>],
    dim: usize,
) -> Vec<T>
where
    DefaultAllocator: Allocator<D>,
{
    samples.iter().map(|s| s[dim].clone()).collect()
}

// ============================================================================
// CONSTANT DISTRIBUTION TESTS
// ============================================================================

mod constant {
    use super::*;

    #[test]
    fn test_constructor_valid() {
        let dist = ConstantDensity::new(5.0);
        assert_eq!(dist.constant(), 5.0);
    }

    #[test]
    fn test_density_at_constant() {
        let dist = ConstantDensity::new(2.5);
        assert_density_correct(&dist, 2.5, true);
    }

    #[test]
    fn test_density_outside_constant() {
        let dist = ConstantDensity::new(2.5);
        assert_density_correct(&dist, 2.6, false);
    }

    #[test]
    fn test_sampling_constant_value() {
        let dist = ConstantDensity::new(3.0);
        let samples = collect_samples(&dist, 100, RNG_SEED);

        for sample in samples {
            assert_eq!(sample, 3.0);
        }
    }

    #[test]
    fn test_sampling_determinism() {
        let dist = ConstantDensity::new(7.0);
        assert_sampling_determinism(&dist, RNG_SEED, 50);
    }
}

// ============================================================================
// UNIFORM DISTRIBUTION TESTS
// ============================================================================

mod uniform {
    use super::*;

    #[test]
    fn test_constructor_invalid_equal_bounds() {
        let dist = UniformDensity::new(1.0, 1.0);
        assert!(dist.is_none());
    }

    #[test]
    fn test_constructor_invalid_reversed() {
        let dist = UniformDensity::new(1.0, 0.0);
        assert!(dist.is_none());
    }

    #[test]
    fn test_constructor_valid_standard() {
        let _dist = UniformDensity::new(0.0, 1.0).unwrap();
        // Constructor succeeded if unwrap() didn't panic
    }

    #[test]
    fn test_constructor_valid_wide_range() {
        let _dist = UniformDensity::new(-100.0, 100.0).unwrap();
        // Constructor succeeded if unwrap() didn't panic
    }

    #[test]
    fn test_density_at_boundaries() {
        let dist = UniformDensity::new(0.0, 1.0).unwrap();
        // Boundaries are inclusive
        assert_density_correct(&dist, 0.0, true);
        assert_density_correct(&dist, 1.0, true);
    }

    #[test]
    fn test_density_in_domain() {
        let dist = UniformDensity::new(0.0, 1.0).unwrap();
        assert_density_correct(&dist, 0.5, true);
    }

    #[test]
    fn test_density_outside_domain() {
        let dist = UniformDensity::new(0.0, 1.0).unwrap();
        assert_density_correct(&dist, 1.5, false);
    }

    #[test]
    fn test_domain_bounded() {
        let dist = UniformDensity::new(0.0, 1.0).unwrap();
        let domain = dist.domain();
        let sample = SVector::from([0.5]);
        assert!(domain.contains::<U1, U1>(&sample.as_view()));
    }

    #[test]
    fn test_sampling_containment() {
        let dist = UniformDensity::new(0.0, 1.0).unwrap();
        let samples = collect_samples(&dist, N_SAMPLES, RNG_SEED);
        assert_samples_in_bounds(&samples, 0.0, 1.0, true);
    }

    #[test]
    fn test_sampling_determinism() {
        let dist = UniformDensity::new(0.0, 1.0).unwrap();
        assert_sampling_determinism(&dist, RNG_SEED, 100);
    }

    #[test]
    fn test_sampling_mean_convergence() {
        let dist = UniformDensity::new(0.0, 1.0).unwrap();
        let samples = collect_samples(&dist, N_SAMPLES, RNG_SEED);
        let empirical_mean = compute_sample_mean(&samples);
        let theoretical_mean = dist.mean()[0];
        assert_abs_diff_eq!(empirical_mean, theoretical_mean, epsilon = 0.025);
    }

    #[test]
    fn test_sampling_variance_convergence() {
        let dist = UniformDensity::new(0.0, 1.0).unwrap();
        let samples = collect_samples(&dist, N_SAMPLES, RNG_SEED);
        let mean = compute_sample_mean(&samples);
        let empirical_variance = compute_sample_variance(&samples, mean);
        let theoretical_variance = dist.variance()[0];
        assert_abs_diff_eq!(empirical_variance, theoretical_variance, epsilon = 0.025);
    }
}

// ============================================================================
// NORMAL DISTRIBUTION TESTS (UNBOUNDED)
// ============================================================================

mod normal {
    use super::*;

    #[test]
    fn test_constructor_invalid_negative_variance() {
        let dist = NormalDensity::new(0.0, -1.0, None, None);
        assert!(dist.is_none());
    }

    #[test]
    fn test_constructor_invalid_zero_variance() {
        let dist = NormalDensity::new(0.0, 0.0, None, None);
        assert!(dist.is_none());
    }

    #[test]
    fn test_constructor_valid_shifted() {
        let _dist = NormalDensity::new(5.0, 2.0, None, None).unwrap();
        // Constructor succeeded if unwrap() didn't panic
    }

    #[test]
    fn test_constructor_valid_standard() {
        let _dist = NormalDensity::new(0.0, 1.0, None, None).unwrap();
        // Constructor succeeded if unwrap() didn't panic
    }

    #[test]
    fn test_density_at_mean() {
        let dist = NormalDensity::new(0.0, 1.0, None, None).unwrap();
        assert_density_correct(&dist, 0.0, true);
    }

    #[test]
    fn test_density_away_from_mean() {
        let dist = NormalDensity::new(0.0, 1.0, None, None).unwrap();
        assert_density_correct(&dist, 2.0, true);
    }

    #[test]
    fn test_domain_unbounded() {
        let dist = NormalDensity::new(0.0, 1.0, None, None).unwrap();
        let domain = dist.domain();
        // Should accept both small and large values
        let sample_small = SVector::from([-1000.0]);
        let sample_large = SVector::from([1000.0]);
        assert!(domain.contains::<U1, U1>(&sample_small.as_view()));
        assert!(domain.contains::<U1, U1>(&sample_large.as_view()));
    }

    #[test]
    fn test_unbounded_sampling_determinism() {
        let dist = NormalDensity::new(0.0, 1.0, None, None).unwrap();
        assert_sampling_determinism(&dist, RNG_SEED, 100);
    }

    #[test]
    fn test_unbounded_sampling_mean_convergence() {
        let dist = NormalDensity::new(0.0, 1.0, None, None).unwrap();
        let samples = collect_samples(&dist, N_SAMPLES, RNG_SEED);
        let empirical_mean = compute_sample_mean(&samples);
        let theoretical_mean = dist.mean()[0];
        assert_abs_diff_eq!(empirical_mean, theoretical_mean, epsilon = 0.08);
    }

    #[test]
    fn test_unbounded_sampling_variance_convergence() {
        let dist = NormalDensity::new(0.0, 1.0, None, None).unwrap();
        let samples = collect_samples(&dist, N_SAMPLES, RNG_SEED);
        let mean = compute_sample_mean(&samples);
        let empirical_variance = compute_sample_variance(&samples, mean);
        let theoretical_variance = dist.variance()[0];
        assert_abs_diff_eq!(empirical_variance, theoretical_variance, epsilon = 0.08);
    }
}

// ============================================================================
// NORMAL DISTRIBUTION TESTS (BOUNDED)
// ============================================================================

mod normal_bounded {
    use super::*;

    #[test]
    fn test_bounded_sampling_asymmetric_containment() {
        // μ=5, σ=1, bounds=[0, 1] (heavily truncated)
        let dist = NormalDensity::new(5.0, 1.0, Some(0.0), Some(1.0)).unwrap();
        let samples = collect_samples(&dist, N_SAMPLES, RNG_SEED);
        assert_samples_in_bounds(&samples, 0.0, 1.0, true);
    }

    #[test]
    fn test_bounded_sampling_asymmetric_determinism() {
        let dist = NormalDensity::new(5.0, 1.0, Some(0.0), Some(1.0)).unwrap();
        assert_sampling_determinism(&dist, RNG_SEED, 100);
    }

    #[test]
    fn test_bounded_sampling_symmetric_containment() {
        let dist = NormalDensity::new(0.0, 1.0, Some(-3.0), Some(3.0)).unwrap();
        let samples = collect_samples(&dist, N_SAMPLES, RNG_SEED);
        assert_samples_in_bounds(&samples, -3.0, 3.0, true);
    }

    #[test]
    fn test_bounded_sampling_symmetric_determinism() {
        let dist = NormalDensity::new(0.0, 1.0, Some(-3.0), Some(3.0)).unwrap();
        assert_sampling_determinism(&dist, RNG_SEED, 100);
    }

    #[test]
    fn test_bounded_sampling_symmetric_mean() {
        let dist = NormalDensity::new(0.0, 1.0, Some(-3.0), Some(3.0)).unwrap();
        let samples = collect_samples(&dist, N_SAMPLES, RNG_SEED);
        let empirical_mean = compute_sample_mean(&samples);
        let theoretical_mean = dist.mean()[0];
        assert_abs_diff_eq!(empirical_mean, theoretical_mean, epsilon = 0.1);
    }

    #[test]
    fn test_density_outside_bounds() {
        let dist = NormalDensity::new(0.0, 1.0, Some(-2.0), Some(2.0)).unwrap();
        assert_density_correct(&dist, 3.0, false);
    }

    #[test]
    fn test_density_within_bounds() {
        let dist = NormalDensity::new(0.0, 1.0, Some(-2.0), Some(2.0)).unwrap();
        assert_density_correct(&dist, 0.0, true);
    }

    #[test]
    fn test_domain_bounded_inclusive() {
        let dist = NormalDensity::new(0.0, 1.0, Some(-2.0), Some(2.0)).unwrap();
        let domain = dist.domain();
        // Boundaries should be inclusive
        let at_min = SVector::from([-2.0]);
        let at_max = SVector::from([2.0]);
        assert!(domain.contains::<U1, U1>(&at_min.as_view()));
        assert!(domain.contains::<U1, U1>(&at_max.as_view()));
    }
}

// ============================================================================
// COSINE DISTRIBUTION TESTS
// ============================================================================

mod cosine {
    use super::*;

    #[test]
    fn test_constructor_invalid_reversed() {
        let dist = CosineDensity::new(0.2, 0.1);
        assert!(dist.is_none());
    }

    #[test]
    fn test_constructor_valid() {
        let _dist = CosineDensity::new(0.1, 0.2).unwrap();
        // Constructor succeeded if unwrap() didn't panic
    }

    #[test]
    fn test_density_outside_domain() {
        let dist = CosineDensity::new(0.1, 0.2).unwrap();
        assert_density_correct(&dist, 0.0, false);
    }

    #[test]
    fn test_density_within_domain() {
        let dist = CosineDensity::new(0.1, 0.2).unwrap();
        assert_density_correct(&dist, 0.15, true);
    }

    #[test]
    fn test_domain_bounded() {
        let dist = CosineDensity::new(0.1, 0.2).unwrap();
        let domain = dist.domain();
        let sample_in = SVector::from([0.15]);
        let sample_out = SVector::from([0.5]);
        assert!(domain.contains::<U1, U1>(&sample_in.as_view()));
        assert!(!domain.contains::<U1, U1>(&sample_out.as_view()));
    }

    #[test]
    fn test_sampling_containment() {
        let dist = CosineDensity::new(0.1, 0.2).unwrap();
        let samples = collect_samples(&dist, N_SAMPLES, RNG_SEED);
        assert_samples_in_bounds(&samples, 0.1, 0.2, true);
    }

    #[test]
    fn test_sampling_determinism() {
        let dist = CosineDensity::new(0.1, 0.2).unwrap();
        assert_sampling_determinism(&dist, RNG_SEED, 100);
    }
}

// ============================================================================
// LOGNORMAL DISTRIBUTION TESTS
// ============================================================================

mod lognormal {
    use super::*;

    #[test]
    fn test_constructor_invalid_zero_sigma() {
        let dist = LognormalDensity::new(0.0, 0.0, 0.1, 10.0);
        assert!(dist.is_none());
    }

    #[test]
    fn test_constructor_valid() {
        let _dist = LognormalDensity::new(0.0, 1.0, 0.1, 10.0).unwrap();
        // Constructor succeeded if unwrap() didn't panic
    }

    #[test]
    fn test_density_outside_domain() {
        let dist = LognormalDensity::new(0.0, 1.0, 0.1, 10.0).unwrap();
        assert_density_correct(&dist, 0.0, false);
    }

    #[test]
    fn test_density_within_domain() {
        let dist = LognormalDensity::new(0.0, 1.0, 0.1, 10.0).unwrap();
        assert_density_correct(&dist, 1.0, true);
    }

    #[test]
    fn test_domain_positive() {
        let dist = LognormalDensity::new(0.0, 1.0, 0.1, 10.0).unwrap();
        let domain = dist.domain();
        let sample_positive = SVector::from([1.0]);
        let sample_negative = SVector::from([-1.0]);
        assert!(domain.contains::<U1, U1>(&sample_positive.as_view()));
        assert!(!domain.contains::<U1, U1>(&sample_negative.as_view()));
    }

    #[test]
    fn test_sampling_containment() {
        let dist = LognormalDensity::new(0.0, 1.0, 0.1, 10.0).unwrap();
        let samples = collect_samples(&dist, N_SAMPLES, RNG_SEED);
        assert_samples_in_bounds(&samples, 0.1, 10.0, true);
    }

    #[test]
    fn test_sampling_determinism() {
        let dist = LognormalDensity::new(0.0, 1.0, 0.1, 10.0).unwrap();
        assert_sampling_determinism(&dist, RNG_SEED, 100);
    }
}

// ============================================================================
// LOG-UNIFORM DISTRIBUTION TESTS
// ============================================================================

mod loguniform {
    use super::*;

    #[test]
    fn test_constructor_invalid_reversed() {
        let dist = LogUniformDensity::new(10.0, 0.1);
        assert!(dist.is_none());
    }

    #[test]
    fn test_constructor_valid() {
        let _dist = LogUniformDensity::new(0.1, 10.0).unwrap();
        // Constructor succeeded if unwrap() didn't panic
    }

    #[test]
    fn test_density_outside_domain() {
        let dist = LogUniformDensity::new(0.1, 10.0).unwrap();
        assert_density_correct(&dist, 0.0, false);
    }

    #[test]
    fn test_density_within_domain() {
        let dist = LogUniformDensity::new(0.1, 10.0).unwrap();
        assert_density_correct(&dist, 1.0, true);
    }

    #[test]
    fn test_domain_positive() {
        let dist = LogUniformDensity::new(0.1, 10.0).unwrap();
        let domain = dist.domain();
        let sample_positive = SVector::from([1.0]);
        let sample_negative = SVector::from([-1.0]);
        assert!(domain.contains::<U1, U1>(&sample_positive.as_view()));
        assert!(!domain.contains::<U1, U1>(&sample_negative.as_view()));
    }

    #[test]
    fn test_sampling_containment() {
        let dist = LogUniformDensity::new(0.1, 10.0).unwrap();
        let samples = collect_samples(&dist, N_SAMPLES, RNG_SEED);
        assert_samples_in_bounds(&samples, 0.1, 10.0, true);
    }

    #[test]
    fn test_sampling_determinism() {
        let dist = LogUniformDensity::new(0.1, 10.0).unwrap();
        assert_sampling_determinism(&dist, RNG_SEED, 100);
    }
}

// ============================================================================
// MULTIVARIATE DISTRIBUTION TESTS
// ============================================================================

mod multivariate {
    use super::*;

    #[test]
    fn test_multivariate_constructor_valid() {
        let univariates = SVector::from([
            UniformDensity::new(0.0, 1.0).unwrap().into(),
            NormalDensity::new(0.0, 1.0, None, None).unwrap().into(),
        ]);
        let _dist = MultivariateDensity::new(univariates);
    }

    #[test]
    fn test_multivariate_sampling_containment() {
        let univariates = SVector::from([
            UniformDensity::new(0.0, 1.0).unwrap().into(),
            UniformDensity::new(-1.0, 1.0).unwrap().into(),
        ]);
        let dist = MultivariateDensity::new(univariates);
        let samples = collect_samples_nd(&dist, N_SAMPLES, RNG_SEED);

        for sample in &samples {
            assert!(sample[0] >= 0.0 && sample[0] <= 1.0, "Dim 0 out of bounds");
            assert!(sample[1] >= -1.0 && sample[1] <= 1.0, "Dim 1 out of bounds");
        }
    }

    #[test]
    fn test_multivariate_sampling_determinism() {
        let univariates = SVector::from([
            UniformDensity::new(0.0, 1.0).unwrap().into(),
            UniformDensity::new(0.0, 1.0).unwrap().into(),
        ]);
        let dist = MultivariateDensity::new(univariates);

        let samples1 = collect_samples_nd(&dist, 50, RNG_SEED);
        let samples2 = collect_samples_nd(&dist, 50, RNG_SEED);

        assert_eq!(
            samples1, samples2,
            "Sampling not deterministic with same seed"
        );
    }

    #[test]
    fn test_multivariate_density_evaluation() {
        let univariates = SVector::from([
            UniformDensity::new(0.0, 1.0).unwrap().into(),
            UniformDensity::new(0.0, 1.0).unwrap().into(),
        ]);
        let dist = MultivariateDensity::new(univariates);

        let sample_in = SVector::from([0.5, 0.5]);
        let sample_out = SVector::from([1.5, 0.5]);

        let density_in = dist.density::<U1, U2>(&sample_in.as_view());
        let density_out = dist.density::<U1, U2>(&sample_out.as_view());

        assert!(density_in.is_some(), "Density should exist inside domain");
        assert!(
            density_out.is_none(),
            "Density should not exist outside domain"
        );
    }

    #[test]
    fn test_multivariate_sampling_mean_convergence() {
        let univariates = SVector::from([
            UniformDensity::new(0.0, 1.0).unwrap().into(),
            UniformDensity::new(1.0, 2.0).unwrap().into(),
        ]);
        let dist = MultivariateDensity::new(univariates);

        let samples = collect_samples_nd(&dist, N_SAMPLES, RNG_SEED);

        // Extract per-dimension samples
        let dim0_samples: Vec<f64> = samples.iter().map(|s| s[0]).collect();
        let dim1_samples: Vec<f64> = samples.iter().map(|s| s[1]).collect();

        let empirical_mean0 = compute_sample_mean(&dim0_samples);
        let empirical_mean1 = compute_sample_mean(&dim1_samples);

        // Theoretical means: Uniform(0,1) has mean 0.5, Uniform(1,2) has mean 1.5
        assert_abs_diff_eq!(empirical_mean0, 0.5, epsilon = 0.025);
        assert_abs_diff_eq!(empirical_mean1, 1.5, epsilon = 0.025);
    }

    #[test]
    fn test_multivariate_sampling_variance_convergence() {
        let univariates = SVector::from([
            NormalDensity::new(0.0, 1.0, None, None).unwrap().into(),
            NormalDensity::new(0.0, 2.0, None, None).unwrap().into(),
        ]);
        let dist = MultivariateDensity::new(univariates);

        let samples = collect_samples_nd(&dist, N_SAMPLES, RNG_SEED);

        // Extract per-dimension samples
        let dim0_samples: Vec<f64> = samples.iter().map(|s| s[0]).collect();
        let dim1_samples: Vec<f64> = samples.iter().map(|s| s[1]).collect();

        let dim0_mean = compute_sample_mean(&dim0_samples);
        let dim1_mean = compute_sample_mean(&dim1_samples);

        let empirical_var0 = compute_sample_variance(&dim0_samples, dim0_mean);
        let empirical_var1 = compute_sample_variance(&dim1_samples, dim1_mean);

        // Theoretical variances: Normal(0,1) has variance 1.0, Normal(0,2) has variance 4.0 (sigma^2)
        assert_abs_diff_eq!(empirical_var0, 1.0, epsilon = 0.1);
        assert_abs_diff_eq!(empirical_var1, 4.0, epsilon = 0.25);
    }
}

// ============================================================================
// MULTIVARIATE NORMAL DISTRIBUTION TESTS
// ============================================================================

mod multinormal {
    use super::*;

    #[test]
    fn test_multinormal_constructor_identity_covariance() {
        let mean = SVector::from([0.0, 0.0]);
        let cov = Matrix2::identity();
        let domain = Domain::new_udomain(U2);
        let _dist = MultivariateNormalDensity::new(cov, domain, Some(mean)).unwrap();
    }

    #[test]
    fn test_multinormal_sampling_containment() {
        let mean = SVector::from([0.0, 0.0]);
        let cov = Matrix2::identity();
        let domain = Domain::new_udomain(U2);
        let dist = MultivariateNormalDensity::new(cov, domain, Some(mean)).unwrap();

        let mut rng = Xoshiro256PlusPlus::seed_from_u64(RNG_SEED);
        let mode = SamplingMode::default();
        let samples: Vec<_> = (0..N_SAMPLES)
            .filter_map(|_| dist.sample(&mut rng, &mode))
            .collect();

        // Normal distribution is unbounded, so just check we got samples
        assert_eq!(samples.len(), N_SAMPLES);
    }

    #[test]
    fn test_multinormal_sampling_determinism() {
        let mean = SVector::from([0.0, 0.0]);
        let cov = Matrix2::identity();
        let domain = Domain::new_udomain(U2);
        let dist = MultivariateNormalDensity::new(cov, domain, Some(mean)).unwrap();

        let mut rng1 = Xoshiro256PlusPlus::seed_from_u64(RNG_SEED);
        let mode = SamplingMode::default();
        let samples1: Vec<_> = (0..50)
            .filter_map(|_| dist.sample(&mut rng1, &mode))
            .collect();

        let mut rng2 = Xoshiro256PlusPlus::seed_from_u64(RNG_SEED);
        let samples2: Vec<_> = (0..50)
            .filter_map(|_| dist.sample(&mut rng2, &mode))
            .collect();

        assert_eq!(
            samples1, samples2,
            "Sampling not deterministic with same seed"
        );
    }

    #[test]
    fn test_multinormal_density_at_mean() {
        let mean = SVector::from([0.0, 0.0]);
        let cov = Matrix2::identity();
        let domain = Domain::new_udomain(U2);
        let dist = MultivariateNormalDensity::new(cov, domain, Some(mean)).unwrap();

        let sample = SVector::from([0.0, 0.0]);
        let density = dist.density::<U1, U2>(&sample.as_view());

        assert!(density.is_some(), "Density should exist at mean");
        assert!(density.unwrap() > 0.0, "Density at mean should be positive");
    }

    #[test]
    fn test_multinormal_density_shifted_mean() {
        let mean = SVector::from([5.0, -3.0]);
        let cov = Matrix2::identity();
        let domain = Domain::new_udomain(U2);
        let dist = MultivariateNormalDensity::new(cov, domain, Some(mean)).unwrap();

        // Density at mean should exist
        let density_at_mean = dist.density::<U1, U2>(&mean.as_view());
        assert!(density_at_mean.is_some(), "Density should exist at mean");

        // Density away from mean should still exist (unbounded)
        let sample_away = SVector::from([0.0, 0.0]);
        let density_away = dist.density::<U1, U2>(&sample_away.as_view());
        assert!(
            density_away.is_some(),
            "Density should exist away from mean"
        );
    }

    #[test]
    fn test_multinormal_sampling_mean_convergence() {
        let mean = SVector::from([2.0, -1.0]);
        let cov = Matrix2::identity();
        let domain = Domain::new_udomain(U2);
        let dist = MultivariateNormalDensity::new(cov, domain, Some(mean)).unwrap();

        let samples = collect_samples_nd(&dist, N_SAMPLES, RNG_SEED);

        // Extract per-dimension samples
        let dim0_samples: Vec<f64> = samples.iter().map(|s| s[0]).collect();
        let dim1_samples: Vec<f64> = samples.iter().map(|s| s[1]).collect();

        let empirical_mean0 = compute_sample_mean(&dim0_samples);
        let empirical_mean1 = compute_sample_mean(&dim1_samples);

        // Theoretical means match the specified mean vector
        assert_abs_diff_eq!(empirical_mean0, 2.0, epsilon = 0.05);
        assert_abs_diff_eq!(empirical_mean1, -1.0, epsilon = 0.05);
    }

    #[test]
    fn test_multinormal_sampling_variance_convergence() {
        let mean = SVector::from([0.0, 0.0]);
        let cov = Matrix2::identity();
        let domain = Domain::new_udomain(U2);
        let dist = MultivariateNormalDensity::new(cov, domain, Some(mean)).unwrap();

        let samples = collect_samples_nd(&dist, N_SAMPLES, RNG_SEED);

        // Extract per-dimension samples
        let dim0_samples: Vec<f64> = samples.iter().map(|s| s[0]).collect();
        let dim1_samples: Vec<f64> = samples.iter().map(|s| s[1]).collect();

        let dim0_mean = compute_sample_mean(&dim0_samples);
        let dim1_mean = compute_sample_mean(&dim1_samples);

        let empirical_var0 = compute_sample_variance(&dim0_samples, dim0_mean);
        let empirical_var1 = compute_sample_variance(&dim1_samples, dim1_mean);

        // Theoretical variances: Identity covariance matrix has diagonal elements of 1.0
        assert_abs_diff_eq!(empirical_var0, 1.0, epsilon = 0.1);
        assert_abs_diff_eq!(empirical_var1, 1.0, epsilon = 0.1);
    }
}

// ============================================================================
// PARTICLE DENSITY DISTRIBUTION TESTS
// ============================================================================

mod particle {
    use super::*;

    #[test]
    fn test_particle_constructor_valid() {
        use nalgebra::OMatrix;

        let mut particles = OMatrix::<f64, U2, Dyn>::zeros(3);
        particles[(0, 0)] = 0.0;
        particles[(1, 0)] = 0.0;
        particles[(0, 1)] = 1.0;
        particles[(1, 1)] = 1.0;
        particles[(0, 2)] = -1.0;
        particles[(1, 2)] = -1.0;

        let domain = Domain::new_udomain(U2);
        let _dist =
            ParticleDensity::from_vectors::<Dyn, Dyn>(&particles.as_view(), domain, None, None)
                .unwrap();
    }

    #[test]
    fn test_particle_unweighted_sampling() {
        use nalgebra::OMatrix;

        let mut particles = OMatrix::<f64, U2, Dyn>::zeros(3);
        particles[(0, 0)] = 0.0;
        particles[(1, 0)] = 0.0;
        particles[(0, 1)] = 1.0;
        particles[(1, 1)] = 1.0;
        particles[(0, 2)] = -1.0;
        particles[(1, 2)] = -1.0;

        let domain = Domain::new_udomain(U2);
        let dist =
            ParticleDensity::from_vectors::<Dyn, Dyn>(&particles.as_view(), domain, None, None)
                .unwrap();

        let mut rng = Xoshiro256PlusPlus::seed_from_u64(RNG_SEED);
        let mode = SamplingMode::default();
        let samples: Vec<_> = (0..100)
            .filter_map(|_| dist.sample(&mut rng, &mode))
            .collect();

        assert_eq!(
            samples.len(),
            100,
            "Should generate correct number of samples"
        );
    }

    #[test]
    fn test_particle_sampling_determinism() {
        use nalgebra::OMatrix;

        let mut particles = OMatrix::<f64, U2, Dyn>::zeros(3);
        particles[(0, 0)] = 0.0;
        particles[(1, 0)] = 0.0;
        particles[(0, 1)] = 1.0;
        particles[(1, 1)] = 1.0;
        particles[(0, 2)] = -1.0;
        particles[(1, 2)] = -1.0;

        let domain = Domain::new_udomain(U2);
        let dist =
            ParticleDensity::from_vectors::<Dyn, Dyn>(&particles.as_view(), domain, None, None)
                .unwrap();

        let mut rng1 = Xoshiro256PlusPlus::seed_from_u64(RNG_SEED);
        let mode = SamplingMode::default();
        let samples1: Vec<_> = (0..50)
            .filter_map(|_| dist.sample(&mut rng1, &mode))
            .collect();

        let mut rng2 = Xoshiro256PlusPlus::seed_from_u64(RNG_SEED);
        let samples2: Vec<_> = (0..50)
            .filter_map(|_| dist.sample(&mut rng2, &mode))
            .collect();

        assert_eq!(
            samples1, samples2,
            "Sampling not deterministic with same seed"
        );
    }

    #[test]
    fn test_particle_density_evaluation() {
        use nalgebra::OMatrix;

        let mut particles = OMatrix::<f64, U2, Dyn>::zeros(2);
        particles[(0, 0)] = 0.0;
        particles[(1, 0)] = 0.0;
        particles[(0, 1)] = 1.0;
        particles[(1, 1)] = 1.0;

        let domain = Domain::new_udomain(U2);
        let dist =
            ParticleDensity::from_vectors::<Dyn, Dyn>(&particles.as_view(), domain, None, None)
                .unwrap();

        // Evaluate at particle location
        let sample_at_particle = SVector::from([0.0, 0.0]);
        let density_at = dist.density::<U1, U2>(&sample_at_particle.as_view());

        // Should have non-zero density at particle location
        assert!(
            density_at.is_some(),
            "Density should exist at particle location"
        );
        assert!(
            density_at.unwrap() > 0.0,
            "Density at particle location should be positive"
        );
    }

    #[test]
    fn test_particle_sampling_mean_convergence() {
        use nalgebra::OMatrix;

        // Create particles centered around (0, 0) and (2, 2)
        let mut particles = OMatrix::<f64, U2, Dyn>::zeros(4);
        particles[(0, 0)] = 0.0;
        particles[(1, 0)] = 0.0;
        particles[(0, 1)] = 0.0;
        particles[(1, 1)] = 0.0;
        particles[(0, 2)] = 2.0;
        particles[(1, 2)] = 2.0;
        particles[(0, 3)] = 2.0;
        particles[(1, 3)] = 2.0;

        let domain = Domain::new_udomain(U2);
        let dist =
            ParticleDensity::from_vectors::<Dyn, Dyn>(&particles.as_view(), domain, None, None)
                .unwrap();

        let samples = collect_samples_nd(&dist, N_SAMPLES, RNG_SEED);

        // Extract per-dimension samples
        let dim0_samples: Vec<f64> = samples.iter().map(|s| s[0]).collect();
        let dim1_samples: Vec<f64> = samples.iter().map(|s| s[1]).collect();

        let empirical_mean0 = compute_sample_mean(&dim0_samples);
        let empirical_mean1 = compute_sample_mean(&dim1_samples);

        // Empirical mean should be close to mean of particles: (0+0+2+2)/4 = 1.0
        assert_abs_diff_eq!(empirical_mean0, 1.0, epsilon = 0.1);
        assert_abs_diff_eq!(empirical_mean1, 1.0, epsilon = 0.1);
    }

    #[test]
    fn test_particle_sampling_variance_convergence() {
        use nalgebra::OMatrix;

        // Create particles at corners of unit square
        let mut particles = OMatrix::<f64, U2, Dyn>::zeros(4);
        particles[(0, 0)] = 0.0;
        particles[(1, 0)] = 0.0;
        particles[(0, 1)] = 1.0;
        particles[(1, 1)] = 0.0;
        particles[(0, 2)] = 0.0;
        particles[(1, 2)] = 1.0;
        particles[(0, 3)] = 1.0;
        particles[(1, 3)] = 1.0;

        let domain = Domain::new_udomain(U2);
        let dist =
            ParticleDensity::from_vectors::<Dyn, Dyn>(&particles.as_view(), domain, None, None)
                .unwrap();

        let samples = collect_samples_nd(&dist, N_SAMPLES, RNG_SEED);

        // Extract per-dimension samples
        let dim0_samples: Vec<f64> = samples.iter().map(|s| s[0]).collect();
        let dim1_samples: Vec<f64> = samples.iter().map(|s| s[1]).collect();

        let dim0_mean = compute_sample_mean(&dim0_samples);
        let dim1_mean = compute_sample_mean(&dim1_samples);

        let empirical_var0 = compute_sample_variance(&dim0_samples, dim0_mean);
        let empirical_var1 = compute_sample_variance(&dim1_samples, dim1_mean);

        // Particles have variance, but kernel smoothing adds uncertainty
        // Just check that variances are reasonable and positive
        assert!(empirical_var0 > 0.0, "Variance should be positive");
        assert!(empirical_var1 > 0.0, "Variance should be positive");
        assert!(
            empirical_var0 < 1.0,
            "Variance should be bounded by particle spread"
        );
        assert!(
            empirical_var1 < 1.0,
            "Variance should be bounded by particle spread"
        );
    }
}

// ============================================================================
// TRAIT CONVERSION TESTS
// ============================================================================
// Tests for newly implemented From trait conversions between distribution types.

#[cfg(test)]
mod trait_conversions {
    use super::*;

    #[test]
    fn test_multivariate_1d_to_univariate_conversion() {
        // Create a 1D MultivariateDensity from a UniformDensity
        let uniform = UniformDensity::new(0.0, 1.0).unwrap();
        let univariate: crate::UnivariateDensity<f64> = uniform.into();
        let marginals = SVector::from([univariate]);
        let multivariate = MultivariateDensity::new(marginals);

        // Convert back to UnivariateDensity using the new From impl
        let recovered: crate::UnivariateDensity<f64> = multivariate.into();

        // Verify the recovered density works
        let sample = SVector::from([0.5]);
        let density = recovered.density::<U1, U1>(&sample.as_view());
        assert!(density.is_some(), "Density should evaluate successfully");
        assert_abs_diff_eq!(density.unwrap(), 1.0, epsilon = 1e-10);
    }

    #[test]
    fn test_multivariate_1d_normal_to_univariate() {
        // Create a 1D MultivariateDensity from a NormalDensity
        let normal = NormalDensity::new(0.0, 1.0, None, None).unwrap();
        let univariate: crate::UnivariateDensity<f64> = normal.into();
        let marginals = SVector::from([univariate]);
        let multivariate = MultivariateDensity::new(marginals);

        // Convert back
        let recovered: crate::UnivariateDensity<f64> = multivariate.into();

        // Verify domain consistency by checking domain bounds
        let domain_recovered = recovered.domain();
        // Normal is unbounded, so max and min should be None
        assert!(domain_recovered.minimum_values()[0].is_none());
        assert!(domain_recovered.maximum_values()[0].is_none());
    }

    #[test]
    fn test_multivariate_1d_constant_to_univariate() {
        // Create a 1D MultivariateDensity from a ConstantDensity
        let constant = ConstantDensity::new(2.5);
        let univariate: crate::UnivariateDensity<f64> = constant.into();
        let marginals = SVector::from([univariate]);
        let multivariate = MultivariateDensity::new(marginals);

        // Convert back
        let recovered: crate::UnivariateDensity<f64> = multivariate.into();

        // Verify the recovered univariate matches constant type
        // Constant values are unbounded in domain
        let sample = SVector::from([0.0]);
        let density = recovered.density::<U1, U1>(&sample.as_view());
        // Constant density should return a value
        assert!(
            density.is_some() || density.is_none(),
            "Test verifies conversion completes"
        );
    }

    #[test]
    fn test_multivariate_1d_preserves_type() {
        // Test that we preserve the type information through the round-trip
        let cosine = CosineDensity::new(0.1, 0.2).unwrap();
        let univariate: crate::UnivariateDensity<f64> = cosine.clone().into();
        let marginals = SVector::from([univariate]);
        let multivariate = MultivariateDensity::new(marginals);

        // Convert back
        let recovered: crate::UnivariateDensity<f64> = multivariate.into();

        // Verify it's a cosine (by checking the typename would work, but we test domain)
        let domain_recovered = recovered.domain();
        assert_abs_diff_eq!(
            domain_recovered.minimum_values()[0].unwrap(),
            0.1,
            epsilon = 1e-10
        );
        assert_abs_diff_eq!(
            domain_recovered.maximum_values()[0].unwrap(),
            0.2,
            epsilon = 1e-10
        );
    }

    #[test]
    fn test_univariate_tryfrom_constant() {
        // Test TryFrom for ConstantDensity
        let constant = ConstantDensity::new(5.0);
        let univariate: crate::UnivariateDensity<f64> = constant.clone().into();

        // Convert back using TryFrom
        let result: Result<ConstantDensity<f64>, _> = univariate.try_into();
        assert!(result.is_ok(), "TryFrom should succeed for ConstantDensity");

        let recovered = result.unwrap();
        assert_abs_diff_eq!(recovered.constant(), 5.0, epsilon = 1e-10);
    }

    #[test]
    fn test_univariate_tryfrom_uniform() {
        // Test TryFrom for UniformDensity
        let uniform = UniformDensity::new(1.0, 3.0).unwrap();
        let univariate: crate::UnivariateDensity<f64> = uniform.clone().into();

        // Convert back using TryFrom
        let result: Result<UniformDensity<f64>, _> = univariate.try_into();
        assert!(result.is_ok(), "TryFrom should succeed for UniformDensity");

        let recovered = result.unwrap();
        assert_abs_diff_eq!(recovered.minimum(), 1.0, epsilon = 1e-10);
        assert_abs_diff_eq!(recovered.maximum(), 3.0, epsilon = 1e-10);
    }

    #[test]
    fn test_univariate_tryfrom_normal() {
        // Test TryFrom for NormalDensity
        let normal = NormalDensity::new(2.0, 0.5, None, None).unwrap();
        let univariate: crate::UnivariateDensity<f64> = normal.clone().into();

        // Convert back using TryFrom
        let result: Result<NormalDensity<f64>, _> = univariate.try_into();
        assert!(result.is_ok(), "TryFrom should succeed for NormalDensity");

        let recovered = result.unwrap();
        let sample = SVector::from([2.0]);
        let density_orig = normal.density::<U1, U1>(&sample.as_view());
        let density_recovered = recovered.density::<U1, U1>(&sample.as_view());
        assert_abs_diff_eq!(
            density_orig.unwrap(),
            density_recovered.unwrap(),
            epsilon = 1e-10
        );
    }

    #[test]
    fn test_univariate_tryfrom_cosine() {
        // Test TryFrom for CosineDensity
        // Use a simple symmetric range that works with cosine
        let cosine = CosineDensity::new(-1.0, 1.0).unwrap();
        let univariate: crate::UnivariateDensity<f64> = cosine.clone().into();

        // Convert back using TryFrom
        let result: Result<CosineDensity<f64>, _> = univariate.try_into();
        assert!(result.is_ok(), "TryFrom should succeed for CosineDensity");

        let recovered = result.unwrap();
        assert_abs_diff_eq!(recovered.minimum(), -1.0, epsilon = 1e-10);
        assert_abs_diff_eq!(recovered.maximum(), 1.0, epsilon = 1e-10);
    }

    #[test]
    fn test_univariate_tryfrom_lognormal() {
        // Test TryFrom for LognormalDensity
        let lognormal = LognormalDensity::new(0.0, 1.0, 0.1, 10.0).unwrap();
        let univariate: crate::UnivariateDensity<f64> = lognormal.clone().into();

        // Convert back using TryFrom
        let result: Result<LognormalDensity<f64>, _> = univariate.try_into();
        assert!(
            result.is_ok(),
            "TryFrom should succeed for LognormalDensity"
        );

        let recovered = result.unwrap();
        let sample = SVector::from([1.0]);
        let density_orig = lognormal.density::<U1, U1>(&sample.as_view());
        let density_recovered = recovered.density::<U1, U1>(&sample.as_view());
        assert_abs_diff_eq!(
            density_orig.unwrap(),
            density_recovered.unwrap(),
            epsilon = 1e-10
        );
    }

    #[test]
    fn test_univariate_tryfrom_loguniform() {
        // Test TryFrom for LogUniformDensity
        let loguniform = LogUniformDensity::new(1.0, 10.0).unwrap();
        let univariate: crate::UnivariateDensity<f64> = loguniform.clone().into();

        // Convert back using TryFrom
        let result: Result<LogUniformDensity<f64>, _> = univariate.try_into();
        assert!(
            result.is_ok(),
            "TryFrom should succeed for LogUniformDensity"
        );

        let recovered = result.unwrap();
        assert_abs_diff_eq!(recovered.minimum(), 1.0, epsilon = 1e-10);
        assert_abs_diff_eq!(recovered.maximum(), 10.0, epsilon = 1e-10);
    }

    #[test]
    fn test_univariate_tryfrom_wrong_type() {
        // Test that TryFrom fails with wrong type
        let uniform = UniformDensity::new(0.0, 1.0).unwrap();
        let univariate: crate::UnivariateDensity<f64> = uniform.into();

        // Try to convert to ConstantDensity (wrong type)
        let result: Result<ConstantDensity<f64>, _> = univariate.try_into();
        assert!(result.is_err(), "TryFrom should fail for wrong type");
    }
}