scirs2-core 0.4.3

//! Advanced numerical methods for ultra-high-performance scientific computing
//!
//! This module implements cutting-edge algorithms for variance reduction, adaptive sampling,
//! and multi-level Monte Carlo methods that are essential for modern scientific computing
//! and computational finance applications.
//!
//! # Key Algorithms
//!
//! - **Multi-level Monte Carlo (MLMC)**: Dramatically reduces computational complexity
//! - **Adaptive sampling**: Dynamically adjusts sampling based on variance estimates
//! - **Antithetic variates**: Variance reduction through negatively correlated samples
//! - **Control variates**: Variance reduction using auxiliary functions
//! - **Importance sampling**: Focuses sampling on high-importance regions
//! - **Sequential Monte Carlo**: Advanced particle filtering methods
//! - **Metropolis-Hastings**: MCMC with adaptive acceptance rates
//!
//! # Examples
//!
//! ```rust
//! use scirs2_core::random::advanced_numerical::*;
//! use scirs2_core::random::core::Random;
//! use rand_distr::Uniform;
//!
//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
//! // Define a computation function that returns the required format for MLMC
//! let compute_option_price = |level: usize, samples: usize| -> Result<Vec<f64>, String> {
//!     // Generate samples for this level
//!     let mut rng = scirs2_core::random::seeded_rng(42 + level as u64);
//!     let mut results = Vec::with_capacity(samples);
//!     for _ in 0..samples {
//!         // Simple Black-Scholes approximation based on level
//!         let dt = 1.0 / (2.0_f64.powi(level as i32));
//!         let random_val = rng.sample(Uniform::new(0.0, 1.0).expect("Operation failed"));
//!         let price = 100.0 * (1.0 + 0.05 * dt * random_val);
//!         results.push((price - 100.0).max(0.0)); // Call option payoff
//!     }
//!     Ok(results)
//! };
//!
//! // Define an expensive computation function
//! let expensive_computation = |rng: &mut Random<rand::rngs::StdRng>| -> f64 {
//!     // Monte Carlo integration of a complex function
//!     let samples: Vec<f64> = (0..100).map(|_| rng.sample(Uniform::new(0.0, 1.0).expect("Operation failed"))).collect();
//!     samples.iter().map(|&x| (x * std::f64::consts::PI).sin().powi(2)).sum::<f64>() / samples.len() as f64
//! };
//!
//! // Multi-level Monte Carlo for option pricing
//! let mlmc = MultiLevelMonteCarlo::new(3, 100); // Smaller values for doc test
//! let estimate = mlmc.estimate(compute_option_price)?;
//!
//! // Adaptive sampling with variance tracking
//! let mut adaptive = AdaptiveSampler::new(0.1, 1000); // More relaxed for doc test
//! let result = adaptive.sample_until_convergence(expensive_computation)?;
//! # Ok(())
//! # }
//! ```

use crate::random::{
    core::{seeded_rng, Random},
    distributions::{Beta, MultivariateNormal},
    parallel::{ParallelRng, ThreadLocalRngPool},
};
use ::ndarray::{Array1, Array2};
use rand::{Rng, RngExt};
use rand_distr::{Distribution, Normal, Uniform};
use std::collections::VecDeque;

/// Multi-level Monte Carlo estimator for variance reduction
///
/// MLMC uses a telescoping sum across multiple levels of approximation to achieve
/// better convergence rates than standard Monte Carlo methods. This is particularly
/// useful for stochastic differential equations and option pricing.
#[derive(Debug, Clone)]
pub struct MultiLevelMonteCarlo {
    max_levels: usize,
    base_samples: usize,
    variance_tolerance: f64,
    convergence_factor: f64,
}

impl MultiLevelMonteCarlo {
    /// Create a new MLMC estimator
    pub fn new(max_levels: usize, base_samples: usize) -> Self {
        Self {
            max_levels,
            base_samples,
            variance_tolerance: 1e-6,
            convergence_factor: 2.0,
        }
    }

    /// Set variance tolerance for adaptive level selection
    pub fn with_tolerance(mut self, tolerance: f64) -> Self {
        self.variance_tolerance = tolerance;
        self
    }

    /// Set convergence factor for sample size scaling
    pub fn with_convergence_factor(mut self, factor: f64) -> Self {
        self.convergence_factor = factor;
        self
    }

    /// Estimate using multi-level Monte Carlo
    pub fn estimate<F>(&self, mut level_function: F) -> Result<MLMCResult, String>
    where
        F: FnMut(usize, usize) -> Result<Vec<f64>, String>,
    {
        let mut estimates = Vec::new();
        let mut variances = Vec::new();
        let mut total_samples = 0;

        for level in 0..self.max_levels {
            // Adaptive sample size calculation
            let level_samples = if level == 0 {
                self.base_samples
            } else {
                (self.base_samples as f64 * self.convergence_factor.powi(level as i32)) as usize
            };

            // Compute level estimate
            let samples = level_function(level, level_samples)?;
            let mean = samples.iter().sum::<f64>() / samples.len() as f64;
            let variance = samples.iter().map(|x| (x - mean).powi(2)).sum::<f64>()
                / (samples.len() - 1) as f64;

            estimates.push(mean);
            variances.push(variance);
            total_samples += level_samples;

            // Check convergence criteria
            if variance < self.variance_tolerance && level > 2 {
                break;
            }
        }

        // Telescoping sum calculation
        let mut mlmc_estimate = estimates[0];
        for i in 1..estimates.len() {
            mlmc_estimate += estimates[i] - estimates[i - 1];
        }

        Ok(MLMCResult {
            estimate: mlmc_estimate,
            variance: variances.iter().sum::<f64>() / variances.len() as f64,
            levels_used: estimates.len(),
            total_samples,
            level_estimates: estimates,
            level_variances: variances,
        })
    }

    /// Parallel MLMC estimation using thread pool
    pub fn estimate_parallel<F>(&self, level_function: F) -> Result<MLMCResult, String>
    where
        F: Fn(usize, usize) -> Result<Vec<f64>, String> + Send + Sync,
    {
        let pool = ThreadLocalRngPool::new(42);

        // Create level tasks
        let level_tasks: Vec<_> = (0..self.max_levels)
            .map(|level| {
                let level_samples = if level == 0 {
                    self.base_samples
                } else {
                    (self.base_samples as f64 * self.convergence_factor.powi(level as i32)) as usize
                };
                (level, level_samples)
            })
            .collect();

        // Execute levels in parallel (simplified implementation)
        let mut level_results = Vec::new();
        for &(level, samples) in &level_tasks {
            let result = level_function(level, samples)?;
            level_results.push(result);
        }

        // Process results
        let mut estimates = Vec::new();
        let mut variances = Vec::new();
        let mut total_samples = 0;

        for (i, samples) in level_results.iter().enumerate() {
            let mean = samples.iter().sum::<f64>() / samples.len() as f64;
            let variance = samples.iter().map(|x| (x - mean).powi(2)).sum::<f64>()
                / (samples.len() - 1) as f64;

            estimates.push(mean);
            variances.push(variance);
            total_samples += samples.len();
        }

        // Telescoping sum
        let mut mlmc_estimate = estimates[0];
        for i in 1..estimates.len() {
            mlmc_estimate += estimates[i] - estimates[i - 1];
        }

        Ok(MLMCResult {
            estimate: mlmc_estimate,
            variance: variances.iter().sum::<f64>() / variances.len() as f64,
            levels_used: estimates.len(),
            total_samples,
            level_estimates: estimates,
            level_variances: variances,
        })
    }
}

/// Result from Multi-level Monte Carlo estimation
#[derive(Debug, Clone)]
pub struct MLMCResult {
    pub estimate: f64,
    pub variance: f64,
    pub levels_used: usize,
    pub total_samples: usize,
    pub level_estimates: Vec<f64>,
    pub level_variances: Vec<f64>,
}

impl MLMCResult {
    /// Calculate confidence interval
    pub fn confidence_interval(&self, confidence: f64) -> (f64, f64) {
        let z_score = match confidence {
            0.90 => 1.645,
            0.95 => 1.96,
            0.99 => 2.576,
            _ => 1.96, // Default to 95%
        };

        let std_error = (self.variance / self.total_samples as f64).sqrt();
        let margin = z_score * std_error;

        (self.estimate - margin, self.estimate + margin)
    }

    /// Calculate relative error
    pub fn relative_error(&self) -> f64 {
        if self.estimate.abs() > 1e-10 {
            (self.variance / self.total_samples as f64).sqrt() / self.estimate.abs()
        } else {
            f64::INFINITY
        }
    }
}

/// Adaptive sampling with dynamic variance tracking
///
/// This sampler automatically adjusts sample sizes based on running variance estimates
/// to achieve desired accuracy with minimal computational cost.
#[derive(Debug)]
pub struct AdaptiveSampler {
    target_tolerance: f64,
    max_samples: usize,
    min_batch_size: usize,
    max_batch_size: usize,
    variance_window: usize,
    running_estimates: VecDeque<f64>,
    running_variances: VecDeque<f64>,
}

impl AdaptiveSampler {
    /// Create a new adaptive sampler
    pub fn new(target_tolerance: f64, max_samples: usize) -> Self {
        Self {
            target_tolerance,
            max_samples,
            min_batch_size: 100,
            max_batch_size: 10000,
            variance_window: 20,
            running_estimates: VecDeque::new(),
            running_variances: VecDeque::new(),
        }
    }

    /// Configure batch size limits
    pub fn with_batch_limits(mut self, min_batch: usize, max_batch: usize) -> Self {
        self.min_batch_size = min_batch;
        self.max_batch_size = max_batch;
        self
    }

    /// Sample until convergence or max samples reached
    pub fn sample_until_convergence<F>(&mut self, mut sampler: F) -> Result<AdaptiveResult, String>
    where
        F: FnMut(&mut Random<rand::rngs::StdRng>) -> f64,
    {
        let mut rng = seeded_rng(42);
        let mut total_samples = 0;
        let mut current_estimate = 0.0;
        let mut current_variance = f64::INFINITY;
        let mut batch_size = self.min_batch_size;

        while total_samples < self.max_samples {
            // Sample current batch
            let mut batch_samples = Vec::with_capacity(batch_size);
            for _ in 0..batch_size {
                batch_samples.push(sampler(&mut rng));
            }

            // Update running statistics
            let batch_mean = batch_samples.iter().sum::<f64>() / batch_samples.len() as f64;
            let batch_variance = batch_samples
                .iter()
                .map(|x| (x - batch_mean).powi(2))
                .sum::<f64>()
                / (batch_samples.len() - 1) as f64;

            self.running_estimates.push_back(batch_mean);
            self.running_variances.push_back(batch_variance);

            // Maintain window size
            if self.running_estimates.len() > self.variance_window {
                self.running_estimates.pop_front();
                self.running_variances.pop_front();
            }

            // Update global estimates
            let total_weight = self.running_estimates.len() as f64;
            current_estimate = self.running_estimates.iter().sum::<f64>() / total_weight;
            current_variance = self.running_variances.iter().sum::<f64>() / total_weight;

            total_samples += batch_size;

            // Check convergence
            let std_error = (current_variance / total_samples as f64).sqrt();
            let relative_error = if current_estimate.abs() > 1e-10 {
                std_error / current_estimate.abs()
            } else {
                std_error
            };

            if relative_error < self.target_tolerance {
                break;
            }

            // Adaptive batch size adjustment
            batch_size = self.adapt_batch_size(current_variance, total_samples, relative_error);
        }

        Ok(AdaptiveResult {
            estimate: current_estimate,
            variance: current_variance,
            samples_used: total_samples,
            converged: self.check_convergence(current_estimate, current_variance, total_samples),
            final_batch_size: batch_size,
        })
    }

    /// Adapt batch size based on current variance and convergence rate
    fn adapt_batch_size(&self, variance: f64, samples_so_far: usize, relative_error: f64) -> usize {
        // Increase batch size if variance is high or we're far from target
        let variance_factor = (variance / self.target_tolerance).sqrt().max(0.1).min(10.0);
        let error_factor = (relative_error / self.target_tolerance).max(0.1).min(10.0);

        let suggested_size = (self.min_batch_size as f64 * variance_factor * error_factor) as usize;
        suggested_size
            .max(self.min_batch_size)
            .min(self.max_batch_size)
    }

    /// Check if sampling has converged
    fn check_convergence(&self, estimate: f64, variance: f64, total_samples: usize) -> bool {
        let std_error = (variance / total_samples as f64).sqrt();
        let relative_error = if estimate.abs() > 1e-10 {
            std_error / estimate.abs()
        } else {
            std_error
        };

        relative_error < self.target_tolerance
    }
}

/// Result from adaptive sampling
#[derive(Debug, Clone)]
pub struct AdaptiveResult {
    pub estimate: f64,
    pub variance: f64,
    pub samples_used: usize,
    pub converged: bool,
    pub final_batch_size: usize,
}

impl AdaptiveResult {
    /// Calculate confidence interval
    pub fn confidence_interval(&self, confidence: f64) -> (f64, f64) {
        let z_score = match confidence {
            0.90 => 1.645,
            0.95 => 1.96,
            0.99 => 2.576,
            _ => 1.96,
        };

        let std_error = (self.variance / self.samples_used as f64).sqrt();
        let margin = z_score * std_error;

        (self.estimate - margin, self.estimate + margin)
    }
}

/// Importance sampling for focusing on high-importance regions
pub struct ImportanceSampler {
    proposal_distribution: Box<dyn Fn(&mut Random<rand::rngs::StdRng>) -> f64 + Send + Sync>,
    target_density: Box<dyn Fn(f64) -> f64 + Send + Sync>,
    proposal_density: Box<dyn Fn(f64) -> f64 + Send + Sync>,
}

impl ImportanceSampler {
    /// Create importance sampler with custom proposal distribution
    pub fn new<P, T, Q>(proposal: P, target_density: T, proposal_density: Q) -> Self
    where
        P: Fn(&mut Random<rand::rngs::StdRng>) -> f64 + Send + Sync + 'static,
        T: Fn(f64) -> f64 + Send + Sync + 'static,
        Q: Fn(f64) -> f64 + Send + Sync + 'static,
    {
        Self {
            proposal_distribution: Box::new(proposal),
            target_density: Box::new(target_density),
            proposal_density: Box::new(proposal_density),
        }
    }

    /// Estimate integral using importance sampling
    pub fn estimate<F>(&self, function: F, num_samples: usize) -> Result<ImportanceResult, String>
    where
        F: Fn(f64) -> f64,
    {
        let mut rng = seeded_rng(42);
        let mut weighted_sum = 0.0;
        let mut weight_sum = 0.0;
        let mut weights = Vec::with_capacity(num_samples);

        for _ in 0..num_samples {
            // Sample from proposal distribution
            let x = (self.proposal_distribution)(&mut rng);

            // Calculate importance weight
            let target_val = (self.target_density)(x);
            let proposal_val = (self.proposal_density)(x);

            if proposal_val > 1e-10 {
                let weight = target_val / proposal_val;
                let function_val = function(x);

                weighted_sum += weight * function_val;
                weight_sum += weight;
                weights.push(weight);
            }
        }

        // Calculate effective sample size
        let weight_sum_sq = weights.iter().map(|w| w * w).sum::<f64>();
        let effective_sample_size = weight_sum * weight_sum / weight_sum_sq;

        let estimate = if weight_sum > 1e-10 {
            weighted_sum / weight_sum
        } else {
            return Err("Zero weight sum in importance sampling".to_string());
        };

        // Estimate variance
        let mut variance_sum = 0.0;
        let mut weight_index = 0;
        let mut rng2 = seeded_rng(42); // Reset for consistent sampling

        for _ in 0..num_samples {
            let x = (self.proposal_distribution)(&mut rng2);
            let target_val = (self.target_density)(x);
            let proposal_val = (self.proposal_density)(x);

            if proposal_val > 1e-10 && weight_index < weights.len() {
                let weight = weights[weight_index];
                let function_val = function(x);
                let weighted_val = weight * function_val / weight_sum;
                variance_sum += weight * (function_val - estimate).powi(2);
                weight_index += 1;
            }
        }

        let variance = variance_sum / (weight_sum * (num_samples - 1) as f64);

        Ok(ImportanceResult {
            estimate,
            variance,
            effective_sample_size,
            total_samples: num_samples,
        })
    }
}

/// Result from importance sampling
#[derive(Debug, Clone)]
pub struct ImportanceResult {
    pub estimate: f64,
    pub variance: f64,
    pub effective_sample_size: f64,
    pub total_samples: usize,
}

/// Sequential Monte Carlo (Particle Filter) implementation
#[derive(Debug)]
pub struct SequentialMonteCarlo {
    num_particles: usize,
    resampling_threshold: f64,
    particles: Vec<Particle>,
}

#[derive(Debug, Clone)]
struct Particle {
    state: Vec<f64>,
    weight: f64,
    log_weight: f64,
}

impl SequentialMonteCarlo {
    /// Create new SMC sampler
    pub fn new(num_particles: usize) -> Self {
        Self {
            num_particles,
            resampling_threshold: 0.5,
            particles: Vec::with_capacity(num_particles),
        }
    }

    /// Initialize particles
    pub fn initialize<F>(&mut self, mut initializer: F) -> Result<(), String>
    where
        F: FnMut(&mut Random<rand::rngs::StdRng>) -> Vec<f64>,
    {
        let mut rng = seeded_rng(42);
        self.particles.clear();

        for _ in 0..self.num_particles {
            let state = initializer(&mut rng);
            self.particles.push(Particle {
                state,
                weight: 1.0 / self.num_particles as f64,
                log_weight: -(self.num_particles as f64).ln(),
            });
        }

        Ok(())
    }

    /// Prediction step
    pub fn predict<F>(&mut self, mut transition: F) -> Result<(), String>
    where
        F: FnMut(&Vec<f64>, &mut Random<rand::rngs::StdRng>) -> Vec<f64>,
    {
        let mut rng = seeded_rng(42);

        for particle in &mut self.particles {
            particle.state = transition(&particle.state, &mut rng);
        }

        Ok(())
    }

    /// Update step with observations
    pub fn update<F>(&mut self, observation: &[f64], mut likelihood: F) -> Result<(), String>
    where
        F: FnMut(&Vec<f64>, &[f64]) -> f64,
    {
        // Update weights based on likelihood
        let mut max_log_weight = f64::NEG_INFINITY;

        for particle in &mut self.particles {
            let likelihood_val = likelihood(&particle.state, observation);
            particle.log_weight += likelihood_val.ln();
            max_log_weight = max_log_weight.max(particle.log_weight);
        }

        // Normalize weights (log-sum-exp trick)
        let mut weight_sum = 0.0;
        for particle in &mut self.particles {
            particle.weight = (particle.log_weight - max_log_weight).exp();
            weight_sum += particle.weight;
        }

        for particle in &mut self.particles {
            particle.weight /= weight_sum;
        }

        // Check if resampling is needed
        let effective_sample_size = self.effective_sample_size();
        if effective_sample_size < self.resampling_threshold * self.num_particles as f64 {
            self.resample()?;
        }

        Ok(())
    }

    /// Calculate effective sample size
    fn effective_sample_size(&self) -> f64 {
        let weight_sum_sq: f64 = self.particles.iter().map(|p| p.weight.powi(2)).sum();
        1.0 / weight_sum_sq
    }

    /// Systematic resampling
    fn resample(&mut self) -> Result<(), String> {
        let mut rng = seeded_rng(42);
        let u0 = rng
            .sample(Uniform::new(0.0, 1.0 / self.num_particles as f64).expect("Operation failed"));

        let mut new_particles = Vec::with_capacity(self.num_particles);
        let mut cumulative_weight = 0.0;
        let mut i = 0;

        for j in 0..self.num_particles {
            let uj = u0 + j as f64 / self.num_particles as f64;

            while cumulative_weight < uj && i < self.particles.len() {
                cumulative_weight += self.particles[i].weight;
                i += 1;
            }

            if i > 0 {
                let mut new_particle = self.particles[i - 1].clone();
                new_particle.weight = 1.0 / self.num_particles as f64;
                new_particle.log_weight = -(self.num_particles as f64).ln();
                new_particles.push(new_particle);
            }
        }

        self.particles = new_particles;
        Ok(())
    }

    /// Get current state estimate (weighted mean)
    pub fn state_estimate(&self) -> Vec<f64> {
        if self.particles.is_empty() {
            return Vec::new();
        }

        let state_dim = self.particles[0].state.len();
        let mut estimate = vec![0.0; state_dim];

        for particle in &self.particles {
            for (i, &val) in particle.state.iter().enumerate() {
                estimate[i] += particle.weight * val;
            }
        }

        estimate
    }

    /// Get covariance matrix of current state
    pub fn state_covariance(&self) -> Array2<f64> {
        let estimate = self.state_estimate();
        let state_dim = estimate.len();
        let mut covariance = Array2::zeros((state_dim, state_dim));

        for particle in &self.particles {
            for i in 0..state_dim {
                for j in 0..state_dim {
                    let diff_i = particle.state[i] - estimate[i];
                    let diff_j = particle.state[j] - estimate[j];
                    covariance[[i, j]] += particle.weight * diff_i * diff_j;
                }
            }
        }

        covariance
    }
}

/// Adaptive Metropolis-Hastings sampler with automatic tuning
#[derive(Debug)]
pub struct AdaptiveMetropolisHastings {
    target_acceptance_rate: f64,
    adaptation_rate: f64,
    proposal_covariance: Array2<f64>,
    accepted_samples: usize,
    total_proposals: usize,
    state_history: VecDeque<Vec<f64>>,
    adaptation_window: usize,
}

impl AdaptiveMetropolisHastings {
    /// Create new adaptive MH sampler
    pub fn new(dimension: usize, target_acceptance: f64) -> Self {
        let mut proposal_cov = Array2::eye(dimension);
        proposal_cov *= 0.1; // Small initial proposal variance

        Self {
            target_acceptance_rate: target_acceptance,
            adaptation_rate: 0.01,
            proposal_covariance: proposal_cov,
            accepted_samples: 0,
            total_proposals: 0,
            state_history: VecDeque::new(),
            adaptation_window: 100,
        }
    }

    /// Sample from target distribution
    pub fn sample<F>(
        &mut self,
        log_density: F,
        initial_state: Vec<f64>,
        num_samples: usize,
    ) -> Result<Vec<Vec<f64>>, String>
    where
        F: Fn(&[f64]) -> f64,
    {
        let mut rng = seeded_rng(42);
        let mut current_state = initial_state;
        let mut current_log_density = log_density(&current_state);
        let mut samples = Vec::with_capacity(num_samples);

        // Create multivariate normal for proposals
        let mut mvn = MultivariateNormal::new(
            vec![0.0; current_state.len()],
            self.array_to_vec2d(&self.proposal_covariance),
        )
        .map_err(|e| format!("Failed to create MVN: {}", e))?;

        for i in 0..num_samples {
            // Generate proposal
            let proposal_delta = mvn.sample(&mut rng);
            let proposal_state: Vec<f64> = current_state
                .iter()
                .zip(proposal_delta.iter())
                .map(|(&curr, &delta)| curr + delta)
                .collect();

            // Evaluate proposal
            let proposal_log_density = log_density(&proposal_state);

            // Acceptance probability
            let log_alpha = proposal_log_density - current_log_density;
            let accept = if log_alpha >= 0.0 {
                true
            } else {
                let u: f64 = rng.sample(Uniform::new(0.0, 1.0).expect("Operation failed"));
                u.ln() < log_alpha
            };

            self.total_proposals += 1;

            if accept {
                current_state = proposal_state;
                current_log_density = proposal_log_density;
                self.accepted_samples += 1;
            }

            samples.push(current_state.clone());
            self.state_history.push_back(current_state.clone());

            // Adapt proposal covariance
            if i > 0 && i % 50 == 0 {
                self.adapt_proposal_covariance();
            }

            // Maintain history window
            if self.state_history.len() > self.adaptation_window {
                self.state_history.pop_front();
            }
        }

        Ok(samples)
    }

    /// Adapt proposal covariance based on recent samples
    fn adapt_proposal_covariance(&mut self) {
        if self.state_history.len() < 10 {
            return;
        }

        // Calculate current acceptance rate
        let acceptance_rate = self.accepted_samples as f64 / self.total_proposals as f64;

        // Adapt based on acceptance rate
        let scale_factor = if acceptance_rate > self.target_acceptance_rate {
            1.0 + self.adaptation_rate
        } else {
            1.0 - self.adaptation_rate
        };

        // Scale proposal covariance
        self.proposal_covariance *= scale_factor.powi(2);

        // Update covariance based on sample history
        if self.state_history.len() > 20 {
            let sample_cov = self.calculate_sample_covariance();
            let learning_rate = 0.05;

            for i in 0..self.proposal_covariance.nrows() {
                for j in 0..self.proposal_covariance.ncols() {
                    self.proposal_covariance[[i, j]] = (1.0 - learning_rate)
                        * self.proposal_covariance[[i, j]]
                        + learning_rate * sample_cov[[i, j]];
                }
            }
        }
    }

    /// Calculate sample covariance from history
    fn calculate_sample_covariance(&self) -> Array2<f64> {
        let n = self.state_history.len();
        let dim = self.state_history[0].len();

        // Calculate mean
        let mut mean = vec![0.0; dim];
        for state in &self.state_history {
            for (i, &val) in state.iter().enumerate() {
                mean[i] += val;
            }
        }
        for val in &mut mean {
            *val /= n as f64;
        }

        // Calculate covariance
        let mut cov = Array2::zeros((dim, dim));
        for state in &self.state_history {
            for i in 0..dim {
                for j in 0..dim {
                    let diff_i = state[i] - mean[i];
                    let diff_j = state[j] - mean[j];
                    cov[[i, j]] += diff_i * diff_j;
                }
            }
        }

        cov /= (n - 1) as f64;
        cov
    }

    /// Convert Array2 to Vec<Vec<f64>> for MultivariateNormal
    fn array_to_vec2d(&self, array: &Array2<f64>) -> Vec<Vec<f64>> {
        array.rows().into_iter().map(|row| row.to_vec()).collect()
    }

    /// Get current acceptance rate
    pub fn acceptance_rate(&self) -> f64 {
        if self.total_proposals > 0 {
            self.accepted_samples as f64 / self.total_proposals as f64
        } else {
            0.0
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;

    #[test]
    fn test_mlmc_basic() {
        let mlmc = MultiLevelMonteCarlo::new(3, 100);

        let result = mlmc
            .estimate(|level, samples| {
                // Simple test function: E[X] = 0.5 at all levels
                let mut rng = seeded_rng(42 + level as u64);
                let mut vals = Vec::with_capacity(samples);
                for _ in 0..samples {
                    vals.push(rng.sample(Uniform::new(0.0, 1.0).expect("Operation failed")));
                }
                Ok(vals)
            })
            .expect("Operation failed");

        assert_relative_eq!(result.estimate, 0.5, epsilon = 0.1);
        assert!(result.levels_used > 0);
        assert!(result.total_samples > 0);
    }

    #[test]
    fn test_adaptive_sampler() {
        let mut sampler = AdaptiveSampler::new(0.05, 10000);

        let result = sampler
            .sample_until_convergence(|rng| {
                // Sample from standard normal
                rng.sample(Normal::new(0.0, 1.0).expect("Operation failed"))
            })
            .expect("Operation failed");

        assert_relative_eq!(result.estimate, 0.0, epsilon = 0.1);
        assert!(result.samples_used > 0);
    }

    #[test]
    fn test_importance_sampling() {
        let sampler = ImportanceSampler::new(
            |rng| rng.sample(Normal::new(1.0, 1.0).expect("Operation failed")), // Proposal: N(1,1)
            |x| (-0.5 * x * x).exp(), // Target: N(0,1) density (unnormalized)
            |x| (-0.5 * (x - 1.0).powi(2)).exp(), // Proposal density (unnormalized)
        );

        let result = sampler.estimate(|x| x, 1000).expect("Operation failed");

        // Should estimate E[X] under N(0,1), which is 0
        assert_relative_eq!(result.estimate, 0.0, epsilon = 0.3);
        assert!(result.effective_sample_size > 0.0);
    }

    #[test]
    fn test_sequential_monte_carlo() {
        let mut smc = SequentialMonteCarlo::new(100);

        // Initialize with standard normal
        smc.initialize(|rng| vec![rng.sample(Normal::new(0.0, 1.0).expect("Operation failed"))])
            .expect("Operation failed");

        // Predict step (random walk)
        smc.predict(|state, rng| {
            let noise = rng.sample(Normal::new(0.0, 0.1).expect("Operation failed"));
            vec![state[0] + noise]
        })
        .expect("Operation failed");

        // Update with observation
        let observation = vec![0.5];
        smc.update(&observation, |state, obs| {
            // Gaussian likelihood
            let diff = state[0] - obs[0];
            (-0.5 * diff * diff).exp()
        })
        .expect("Operation failed");

        let estimate = smc.state_estimate();
        assert_eq!(estimate.len(), 1);
        // Should be pulled towards observation
        assert!(estimate[0].abs() < 2.0);
    }

    #[test]
    fn test_adaptive_metropolis_hastings() {
        let mut amh = AdaptiveMetropolisHastings::new(2, 0.44);

        // 10000 samples for reliable convergence; seed=42 is baked into AMH::sample()
        // making this deterministic and not flaky
        let samples = amh
            .sample(
                |state| {
                    // Standard 2D normal log-density
                    -0.5 * (state[0].powi(2) + state[1].powi(2))
                },
                vec![0.0, 0.0],
                10000,
            )
            .expect("Operation failed");

        assert_eq!(samples.len(), 10000);
        assert!(amh.acceptance_rate() > 0.1);
        assert!(amh.acceptance_rate() < 0.9);

        // With 10000 samples and fixed seed=42, empirical mean should be within 0.5
        let mean_x: f64 = samples.iter().map(|s| s[0]).sum::<f64>() / samples.len() as f64;
        let mean_y: f64 = samples.iter().map(|s| s[1]).sum::<f64>() / samples.len() as f64;

        assert_relative_eq!(mean_x, 0.0, epsilon = 0.5);
        assert_relative_eq!(mean_y, 0.0, epsilon = 0.5);
    }
}