//! Stochastic computation graphs for probabilistic programming
//!
//! This module provides support for automatic differentiation through stochastic
//! operations and probabilistic models. It implements various gradient estimators
//! for sampling operations, variance reduction techniques, and tools for building
//! differentiable probabilistic programs.

// Framework infrastructure - components designed for future use
#![allow(dead_code)]
use std::collections::HashMap;
use torsh_core::error::{Result, TorshError};
use torsh_tensor::{creation, Tensor};

/// Configuration for stochastic gradient estimation
#[derive(Debug, Clone)]
pub struct StochasticConfig {
    /// Number of samples for Monte Carlo estimation
    pub num_samples: usize,
    /// Variance reduction technique
    pub variance_reduction: VarianceReduction,
    /// Baseline estimation method
    pub baseline: BaselineEstimator,
    /// Learning rate for baseline updates
    pub baseline_lr: f32,
    /// Control variate coefficient
    pub control_variate_coeff: f32,
    /// Whether to use reparameterization trick when possible
    pub use_reparameterization: bool,
    /// Temperature for Gumbel-based methods
    pub temperature: f32,
}

impl Default for StochasticConfig {
    fn default() -> Self {
        Self {
            num_samples: 100,
            variance_reduction: VarianceReduction::Baseline,
            baseline: BaselineEstimator::MovingAverage,
            baseline_lr: 0.01,
            control_variate_coeff: 1.0,
            use_reparameterization: true,
            temperature: 1.0,
        }
    }
}

/// Variance reduction techniques for stochastic gradients
#[derive(Debug, Clone, PartialEq)]
pub enum VarianceReduction {
    /// No variance reduction
    None,
    /// Baseline subtraction (REINFORCE with baseline)
    Baseline,
    /// Control variates
    ControlVariates,
    /// Importance sampling
    ImportanceSampling,
    /// Multiple importance sampling
    MultipleImportanceSampling,
    /// Rao-Blackwellization
    RaoBlackwell,
    /// Low-variance pathwise derivatives
    PathwiseDerivatives,
}

/// Baseline estimators for variance reduction
#[derive(Debug, Clone, PartialEq)]
pub enum BaselineEstimator {
    /// Moving average baseline
    MovingAverage,
    /// Neural network baseline
    NeuralNetwork,
    /// State-dependent baseline
    StateDependent,
    /// Optimal control variate
    OptimalControlVariate,
}

/// Stochastic gradient estimators
#[derive(Debug, Clone, PartialEq)]
pub enum GradientEstimator {
    /// REINFORCE (score function) estimator
    Reinforce,
    /// Reparameterization trick
    Reparameterization,
    /// Gumbel-Softmax
    GumbelSoftmax,
    /// Straight-through estimator
    StraightThrough,
    /// NVIL (Neural Variational Inference and Learning)
    NVIL,
    /// VIMCO (Variational Inference for Monte Carlo Objectives)
    VIMCO,
}

/// Stochastic operations that can appear in computation graphs
pub trait StochasticOperation {
    /// Sample from the distribution
    fn sample(&self, params: &[&Tensor], _config: &StochasticConfig) -> Result<Tensor>;

    /// Compute log probability of a sample
    fn log_prob(&self, sample: &Tensor, params: &[&Tensor]) -> Result<Tensor>;

    /// Compute gradients using score function estimator
    fn score_function_gradient(
        &self,
        sample: &Tensor,
        params: &[&Tensor],
        downstream_grad: &Tensor,
        _config: &StochasticConfig,
    ) -> Result<Vec<Tensor>>;

    /// Compute gradients using reparameterization trick (if applicable)
    fn reparameterization_gradient(
        &self,
        params: &[&Tensor],
        downstream_grad: &Tensor,
        _config: &StochasticConfig,
    ) -> Result<Option<Vec<Tensor>>>;

    /// Name of the distribution
    fn name(&self) -> &str;

    /// Whether this operation supports reparameterization
    fn supports_reparameterization(&self) -> bool;
}

/// Normal (Gaussian) distribution
pub struct NormalDistribution;

impl StochasticOperation for NormalDistribution {
    fn sample(&self, params: &[&Tensor], _config: &StochasticConfig) -> Result<Tensor> {
        if params.len() != 2 {
            return Err(TorshError::InvalidArgument(
                "Normal distribution requires mean and std parameters".to_string(),
            ));
        }

        let mean = params[0];
        let std = params[1];
        let shape = mean.shape();

        // Sample from standard normal and transform
        let eps = creation::randn::<f32>(shape.dims())?;
        mean.add(&std.mul(&eps)?)
    }

    fn log_prob(&self, sample: &Tensor, params: &[&Tensor]) -> Result<Tensor> {
        let mean = params[0];
        let std = params[1];

        let normalized = sample.sub(mean)?.div(std)?;
        let log_std = std.log()?;
        let log_2pi = creation::tensor_scalar(2.0 * std::f32::consts::PI)?.log()?;

        // -0.5 * (normalized^2 + log(2π) + 2*log(std))
        let _neg_half = creation::tensor_scalar(-0.5)?;
        let squared_norm = normalized.pow_scalar(2.0)?;
        let log_prob_unnorm = squared_norm.add(&log_2pi)?.add(&log_std.mul_scalar(2.0)?)?;

        log_prob_unnorm.mul_scalar(-0.5)
    }

    fn score_function_gradient(
        &self,
        sample: &Tensor,
        params: &[&Tensor],
        downstream_grad: &Tensor,
        _config: &StochasticConfig,
    ) -> Result<Vec<Tensor>> {
        let mean = params[0];
        let std = params[1];

        // Gradient of log p(x|θ) w.r.t. parameters
        let normalized = sample.sub(mean)?.div(std)?;

        // ∂log p/∂μ = (x - μ) / σ²
        let grad_mean = normalized.div(std)?;

        // ∂log p/∂σ = -1/σ + (x - μ)²/σ³
        let neg_inv_std = std.pow_scalar(-1.0)?.mul_scalar(-1.0)?;
        let normalized_squared_over_std = normalized.pow_scalar(2.0)?.div(std)?;
        let grad_std = neg_inv_std.add(&normalized_squared_over_std)?;

        // Score function: ∇_θ log p(x|θ) * downstream_grad
        let downstream_scalar = downstream_grad.item().unwrap_or(1.0);
        let score_mean = grad_mean.mul_scalar(downstream_scalar)?;
        let score_std = grad_std.mul_scalar(downstream_scalar)?;

        Ok(vec![score_mean, score_std])
    }

    fn reparameterization_gradient(
        &self,
        params: &[&Tensor],
        downstream_grad: &Tensor,
        __config: &StochasticConfig,
    ) -> Result<Option<Vec<Tensor>>> {
        let _mean = params[0];
        let std = params[1];

        // For x = μ + σ * ε, where ε ~ N(0,1):
        // ∂x/∂μ = 1, ∂x/∂σ = ε
        // But we need ε for the gradient, which we don't store
        // In practice, this would require saving the noise sample

        let grad_mean = downstream_grad.clone();
        let grad_std = downstream_grad.mul(std)?; // Simplified

        Ok(Some(vec![grad_mean, grad_std]))
    }

    fn name(&self) -> &str {
        "Normal"
    }

    fn supports_reparameterization(&self) -> bool {
        true
    }
}

/// Bernoulli distribution
pub struct BernoulliDistribution;

impl StochasticOperation for BernoulliDistribution {
    fn sample(&self, params: &[&Tensor], __config: &StochasticConfig) -> Result<Tensor> {
        if params.len() != 1 {
            return Err(TorshError::InvalidArgument(
                "Bernoulli distribution requires probability parameter".to_string(),
            ));
        }

        let probs = params[0];
        let uniform = creation::rand::<f32>(probs.shape().dims())?;

        // Create boolean mask where uniform < probs
        // TODO: Replace with proper tensor comparison when available
        let _mask = uniform.clone(); // Placeholder implementation
                                     // Placeholder: return uniform tensor for now
        Ok(uniform)
    }

    fn log_prob(&self, sample: &Tensor, params: &[&Tensor]) -> Result<Tensor> {
        let probs = params[0];
        let eps = 1e-8;

        // log p(x) = x * log(p) + (1-x) * log(1-p)
        let log_p = probs.add_scalar(eps)?.log()?;
        let log_one_minus_p = probs.neg()?.add_scalar(1.0 + eps)?.log()?;

        let term1 = sample.mul(&log_p)?;
        let term2 = sample.neg()?.add_scalar(1.0)?.mul(&log_one_minus_p)?;

        term1.add(&term2)
    }

    fn score_function_gradient(
        &self,
        sample: &Tensor,
        params: &[&Tensor],
        downstream_grad: &Tensor,
        __config: &StochasticConfig,
    ) -> Result<Vec<Tensor>> {
        let probs = params[0];
        let eps = 1e-8;

        // ∂log p/∂p = x/p - (1-x)/(1-p)
        let term1 = sample.div(&probs.add_scalar(eps)?)?;
        let term2 = sample
            .neg()?
            .add_scalar(1.0)?
            .div(&probs.neg()?.add_scalar(1.0 + eps)?)?;
        let grad_log_prob = term1.add(&term2)?;

        let score_grad = grad_log_prob.mul(downstream_grad)?;

        Ok(vec![score_grad])
    }

    fn reparameterization_gradient(
        &self,
        _params: &[&Tensor],
        _downstream_grad: &Tensor,
        __config: &StochasticConfig,
    ) -> Result<Option<Vec<Tensor>>> {
        // Bernoulli doesn't naturally support reparameterization
        // Could use Gumbel-Softmax or straight-through estimator
        Ok(None)
    }

    fn name(&self) -> &str {
        "Bernoulli"
    }

    fn supports_reparameterization(&self) -> bool {
        false
    }
}

/// Categorical distribution with Gumbel-Softmax reparameterization
pub struct CategoricalDistribution;

impl StochasticOperation for CategoricalDistribution {
    fn sample(&self, params: &[&Tensor], config: &StochasticConfig) -> Result<Tensor> {
        if params.len() != 1 {
            return Err(TorshError::InvalidArgument(
                "Categorical distribution requires logits parameter".to_string(),
            ));
        }

        let logits = params[0];

        if config.use_reparameterization {
            // Gumbel-Softmax sampling
            self.gumbel_softmax_sample(logits, config.temperature)
        } else {
            // Standard categorical sampling
            let probs = logits.softmax(-1)?;
            self.categorical_sample(&probs)
        }
    }

    fn log_prob(&self, sample: &Tensor, params: &[&Tensor]) -> Result<Tensor> {
        let logits = params[0];
        let log_probs = logits.log_softmax(-1)?;

        // For one-hot sample: sum(sample * log_probs)
        sample.mul(&log_probs)?.sum_dim(&[-1], false)
    }

    fn score_function_gradient(
        &self,
        sample: &Tensor,
        params: &[&Tensor],
        downstream_grad: &Tensor,
        __config: &StochasticConfig,
    ) -> Result<Vec<Tensor>> {
        let logits = params[0];
        let probs = logits.softmax(-1)?;

        // ∂log p/∂logits = sample - probs
        let grad_log_prob = sample.sub(&probs)?;
        let score_grad = grad_log_prob.mul(downstream_grad)?;

        Ok(vec![score_grad])
    }

    fn reparameterization_gradient(
        &self,
        params: &[&Tensor],
        downstream_grad: &Tensor,
        config: &StochasticConfig,
    ) -> Result<Option<Vec<Tensor>>> {
        let logits = params[0];

        // Gumbel-Softmax gradient
        let gumbel_noise = self.sample_gumbel(&logits.shape())?;
        let noisy_logits = logits.add(&gumbel_noise)?;
        let softmax_grad =
            self.gumbel_softmax_gradient(&noisy_logits, config.temperature, downstream_grad)?;

        Ok(Some(vec![softmax_grad]))
    }

    fn name(&self) -> &str {
        "Categorical"
    }

    fn supports_reparameterization(&self) -> bool {
        true
    }
}

impl CategoricalDistribution {
    fn gumbel_softmax_sample(&self, logits: &Tensor, temperature: f32) -> Result<Tensor> {
        let gumbel_noise = self.sample_gumbel(&logits.shape())?;
        let noisy_logits = logits.add(&gumbel_noise)?;
        let scaled_logits = noisy_logits.div_scalar(temperature)?;
        scaled_logits.softmax(-1)
    }

    fn gumbel_softmax_gradient(
        &self,
        noisy_logits: &Tensor,
        temperature: f32,
        downstream_grad: &Tensor,
    ) -> Result<Tensor> {
        let softmax_output = noisy_logits.div_scalar(temperature)?.softmax(-1)?;

        // Gradient of softmax w.r.t. logits
        let grad = downstream_grad.mul(&softmax_output)?;
        let sum_grad = grad.sum_dim(&[-1], true)?;
        let softmax_grad = grad.sub(&softmax_output.mul(&sum_grad)?)?;

        softmax_grad.div_scalar(temperature)
    }

    fn categorical_sample(&self, probs: &Tensor) -> Result<Tensor> {
        // Sample from categorical distribution using inverse CDF
        let uniform = creation::rand::<f32>(probs.shape().dims())?;
        let cumsum = probs.cumsum(-1)?;

        // Find the first index where cumsum >= uniform
        let _expanded_uniform = uniform.unsqueeze(-1)?;

        // Find the first index where cumsum >= uniform
        // Create boolean comparison where cumsum >= expanded_uniform
        // TODO: Replace with proper tensor comparison when available
        let comparison = cumsum.clone(); // Placeholder implementation

        // Create indices tensor - simplified implementation
        // This is a placeholder that returns zeros for now
        // TODO: Implement proper argmax for finding first True value
        let indices_shape = comparison.shape();
        let indices: Tensor<f32> = creation::zeros(indices_shape.dims())?;
        Ok(indices)
    }

    fn sample_gumbel(&self, shape: &torsh_core::shape::Shape) -> Result<Tensor> {
        let uniform = creation::rand::<f32>(&shape.dims())?;
        let eps = 1e-20;
        let log_uniform = uniform.add_scalar(eps)?.log()?;
        let neg_log_uniform = log_uniform.neg()?;
        neg_log_uniform.add_scalar(eps)?.log()?.neg()
    }
}

/// Stochastic computation graph builder
pub struct StochasticGraph {
    operations: Vec<Box<dyn StochasticOperation>>,
    dependencies: HashMap<usize, Vec<usize>>,
    config: StochasticConfig,
    baseline_values: HashMap<String, f32>,
}

impl StochasticGraph {
    pub fn new(config: StochasticConfig) -> Self {
        Self {
            operations: Vec::new(),
            dependencies: HashMap::new(),
            config,
            baseline_values: HashMap::new(),
        }
    }

    pub fn add_operation<T: StochasticOperation + 'static>(&mut self, op: T) -> usize {
        let id = self.operations.len();
        self.operations.push(Box::new(op));
        id
    }

    pub fn add_dependency(&mut self, node_id: usize, dependency_id: usize) {
        self.dependencies
            .entry(node_id)
            .or_insert_with(Vec::new)
            .push(dependency_id);
    }

    /// Execute forward pass with sampling
    pub fn forward(&self, inputs: &HashMap<usize, Tensor>) -> Result<HashMap<usize, Tensor>> {
        let mut outputs = HashMap::new();
        let execution_order = self.topological_sort()?;

        for node_id in execution_order {
            if let Some(input) = inputs.get(&node_id) {
                outputs.insert(node_id, input.clone());
            } else if let Some(op) = self.operations.get(node_id) {
                let empty_deps = vec![];
                let deps = self.dependencies.get(&node_id).unwrap_or(&empty_deps);
                let dep_tensors: Vec<&Tensor> = deps
                    .iter()
                    .filter_map(|&dep_id| outputs.get(&dep_id))
                    .collect();

                if !dep_tensors.is_empty() {
                    let result = op.sample(&dep_tensors, &self.config)?;
                    outputs.insert(node_id, result);
                }
            }
        }

        Ok(outputs)
    }

    /// Compute gradients using chosen estimator
    pub fn backward(
        &mut self,
        outputs: &HashMap<usize, Tensor>,
        target_loss: &Tensor,
        estimator: GradientEstimator,
    ) -> Result<HashMap<usize, Tensor>> {
        match estimator {
            GradientEstimator::Reinforce => self.reinforce_backward(outputs, target_loss),
            GradientEstimator::Reparameterization => {
                self.reparameterization_backward(outputs, target_loss)
            }
            GradientEstimator::GumbelSoftmax => self.gumbel_softmax_backward(outputs, target_loss),
            _ => Err(TorshError::NotImplemented(
                "Gradient estimator not implemented".to_string(),
            )),
        }
    }

    fn reinforce_backward(
        &mut self,
        outputs: &HashMap<usize, Tensor>,
        target_loss: &Tensor,
    ) -> Result<HashMap<usize, Tensor>> {
        let mut gradients = HashMap::new();
        let baseline = self.compute_baseline(target_loss)?;

        let mut centered_loss = target_loss.clone();
        centered_loss.sub_scalar_(baseline)?;

        for (node_id, sample) in outputs {
            if let Some(op) = self.operations.get(*node_id) {
                let empty_vec = vec![];
                let deps = self.dependencies.get(node_id).unwrap_or(&empty_vec);
                let dep_tensors: Vec<&Tensor> = deps
                    .iter()
                    .filter_map(|&dep_id| outputs.get(&dep_id))
                    .collect();

                if !dep_tensors.is_empty() {
                    let param_grads = op.score_function_gradient(
                        sample,
                        &dep_tensors,
                        &centered_loss,
                        &self.config,
                    )?;

                    for (i, grad) in param_grads.into_iter().enumerate() {
                        gradients.insert(*node_id * 1000 + i, grad);
                    }
                }
            }
        }

        self.update_baseline(target_loss.item()?)?;
        Ok(gradients)
    }

    fn reparameterization_backward(
        &self,
        outputs: &HashMap<usize, Tensor>,
        target_loss: &Tensor,
    ) -> Result<HashMap<usize, Tensor>> {
        let mut gradients = HashMap::new();

        for (node_id, _sample) in outputs {
            if let Some(op) = self.operations.get(*node_id) {
                if op.supports_reparameterization() {
                    let empty_vec = vec![];
                    let deps = self.dependencies.get(node_id).unwrap_or(&empty_vec);
                    let dep_tensors: Vec<&Tensor> = deps
                        .iter()
                        .filter_map(|&dep_id| outputs.get(&dep_id))
                        .collect();

                    if !dep_tensors.is_empty() {
                        if let Some(param_grads) =
                            op.reparameterization_gradient(&dep_tensors, target_loss, &self.config)?
                        {
                            for (i, grad) in param_grads.into_iter().enumerate() {
                                gradients.insert(*node_id * 1000 + i, grad);
                            }
                        }
                    }
                }
            }
        }

        Ok(gradients)
    }

    fn gumbel_softmax_backward(
        &self,
        _outputs: &HashMap<usize, Tensor>,
        _target_loss: &Tensor,
    ) -> Result<HashMap<usize, Tensor>> {
        // Implementation for Gumbel-Softmax backward pass
        Ok(HashMap::new())
    }

    fn compute_baseline(&self, loss: &Tensor) -> Result<f32> {
        match self.config.baseline {
            BaselineEstimator::MovingAverage => {
                let current_loss = loss.to_vec()?[0];
                if let Some(&prev_baseline) = self.baseline_values.get("moving_average") {
                    Ok(prev_baseline * (1.0 - self.config.baseline_lr)
                        + current_loss * self.config.baseline_lr)
                } else {
                    Ok(current_loss)
                }
            }
            _ => Ok(0.0), // Other baseline estimators not implemented
        }
    }

    fn update_baseline(&mut self, loss: f32) -> Result<()> {
        let baseline = self.compute_baseline(&creation::tensor_scalar(loss)?)?;
        self.baseline_values
            .insert("moving_average".to_string(), baseline);
        Ok(())
    }

    fn topological_sort(&self) -> Result<Vec<usize>> {
        // Simple topological sort implementation
        let mut sorted = Vec::new();
        let mut visited = std::collections::HashSet::new();

        for i in 0..self.operations.len() {
            if !visited.contains(&i) {
                self.dfs_visit(i, &mut visited, &mut sorted);
            }
        }

        sorted.reverse();
        Ok(sorted)
    }

    fn dfs_visit(
        &self,
        node: usize,
        visited: &mut std::collections::HashSet<usize>,
        sorted: &mut Vec<usize>,
    ) {
        visited.insert(node);

        if let Some(deps) = self.dependencies.get(&node) {
            for &dep in deps {
                if !visited.contains(&dep) {
                    self.dfs_visit(dep, visited, sorted);
                }
            }
        }

        sorted.push(node);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    // use torsh_core::shape::Shape;

    #[test]
    fn test_normal_distribution() {
        let normal = NormalDistribution;
        let mean = Tensor::zeros(&[5], torsh_core::DeviceType::Cpu).unwrap();
        let std = Tensor::ones(&[5], torsh_core::DeviceType::Cpu).unwrap();
        let params = vec![&mean, &std];

        let config = StochasticConfig::default();
        let sample = normal.sample(&params, &config).unwrap();

        assert_eq!(sample.shape().dims(), &[5]);

        let log_prob = normal.log_prob(&sample, &params).unwrap();
        // log_prob may be scalar if implementation returns total log probability
        let shape = log_prob.shape();
        let dims = shape.dims();
        assert!(
            dims == &[5] || dims == &[] as &[usize],
            "Expected log_prob shape to be [5] or [], got {:?}",
            dims
        );
    }

    #[test]
    fn test_bernoulli_distribution() {
        let bernoulli = BernoulliDistribution;
        let probs = creation::full(&[10], 0.5).unwrap();
        let params = vec![&probs];

        let config = StochasticConfig::default();
        let sample = bernoulli.sample(&params, &config).unwrap();

        assert_eq!(sample.shape().dims(), &[10]);

        let log_prob = bernoulli.log_prob(&sample, &params).unwrap();
        assert_eq!(log_prob.shape().dims(), &[10]);
    }

    #[test]
    fn test_categorical_gumbel_softmax() {
        let categorical = CategoricalDistribution;
        // Use controlled logits with small random values for better behavior
        let logits = torsh_tensor::creation::full::<f32>(&[5, 3], 0.1).unwrap();
        let params = vec![&logits];

        let mut config = StochasticConfig::default();
        config.use_reparameterization = true;
        config.temperature = 1.0; // Standard temperature

        let sample = categorical.sample(&params, &config).unwrap();
        assert_eq!(sample.shape().dims(), &[5, 3]);

        // Check that it's approximately one-hot (with appropriate tolerance for Gumbel-Softmax)
        let sums = sample.sum_dim(&[-1], false).unwrap();
        for i in 0..5 {
            let sum_val = sums.select(0, i).unwrap().to_vec().unwrap()[0];
            // Gumbel-Softmax produces soft samples where sums should be close to 1.0
            // Using reasonable tolerance for numerical stability in continuous relaxation
            assert!(
                (sum_val - 1.0).abs() < 0.5,
                "Sample {} sum {} deviates too much from 1.0",
                i,
                sum_val
            );
        }

        // Additional check: all values should be positive (since it's a softmax)
        let sample_data = sample.to_vec().unwrap();
        for val in sample_data {
            assert!(
                val >= 0.0,
                "Gumbel-Softmax values should be non-negative: {}",
                val
            );
        }
    }

    #[test]
    fn test_score_function_gradient() {
        let normal = NormalDistribution;
        let mean = Tensor::zeros(&[3], torsh_core::DeviceType::Cpu).unwrap();
        let std = Tensor::ones(&[3], torsh_core::DeviceType::Cpu).unwrap();
        let params = vec![&mean, &std];

        let sample = torsh_tensor::creation::randn::<f32>(&[3]).unwrap();
        // Use scalar downstream gradient for score function
        let downstream_grad = torsh_tensor::creation::tensor_scalar(1.0).unwrap();

        let config = StochasticConfig::default();
        let grads = normal
            .score_function_gradient(&sample, &params, &downstream_grad, &config)
            .unwrap();

        assert_eq!(grads.len(), 2); // mean and std gradients
        assert_eq!(grads[0].shape().dims(), &[3]);
        assert_eq!(grads[1].shape().dims(), &[3]);
    }

    #[test]
    fn test_stochastic_graph() {
        let config = StochasticConfig::default();
        let mut graph = StochasticGraph::new(config);

        let normal_id = graph.add_operation(NormalDistribution);
        assert_eq!(normal_id, 0);

        let mut inputs = HashMap::new();
        inputs.insert(0, Tensor::zeros(&[2], torsh_core::DeviceType::Cpu).unwrap());
        inputs.insert(1, Tensor::ones(&[2], torsh_core::DeviceType::Cpu).unwrap());

        graph.add_dependency(normal_id, 0);
        graph.add_dependency(normal_id, 1);

        let outputs = graph.forward(&inputs).unwrap();
        assert!(outputs.contains_key(&normal_id));
    }

    #[test]
    fn test_variance_reduction_config() {
        let config = StochasticConfig {
            variance_reduction: VarianceReduction::ControlVariates,
            baseline: BaselineEstimator::NeuralNetwork,
            num_samples: 200,
            ..Default::default()
        };

        assert_eq!(
            config.variance_reduction,
            VarianceReduction::ControlVariates
        );
        assert_eq!(config.baseline, BaselineEstimator::NeuralNetwork);
        assert_eq!(config.num_samples, 200);
    }

    #[test]
    fn test_gradient_estimators() {
        // Test that different gradient estimators are distinguishable
        let estimators = vec![
            GradientEstimator::Reinforce,
            GradientEstimator::Reparameterization,
            GradientEstimator::GumbelSoftmax,
        ];

        for estimator in estimators {
            match estimator {
                GradientEstimator::Reinforce => assert!(true),
                GradientEstimator::Reparameterization => assert!(true),
                GradientEstimator::GumbelSoftmax => assert!(true),
                _ => assert!(false),
            }
        }
    }
}
