scirs2_core/random/

neural_sampling.rs

//! Neural-based sampling methods for ultra-modern generative modeling
//!
//! This module implements the most advanced neural sampling algorithms from cutting-edge
//! machine learning research. These methods leverage deep neural networks to learn complex
//! probability distributions and generate high-quality samples.
//!
//! # Implemented Methods
//!
//! - **Normalizing Flows**: Invertible neural networks for exact likelihood computation
//! - **Variational Autoencoders (VAE)**: Probabilistic latent variable models
//! - **Score-Based Diffusion Models**: State-of-the-art generative models using score matching
//! - **Energy-Based Models (EBM)**: Flexible unnormalized probability models
//! - **Neural Posterior Estimation**: Amortized Bayesian inference
//! - **Autoregressive Models**: Sequential probability modeling
//! - **Generative Adversarial Sampling**: Adversarial training for sample generation
//!
//! # Key Advantages
//!
//! - **Expressiveness**: Can model highly complex, multi-modal distributions
//! - **Scalability**: Efficient sampling from high-dimensional spaces
//! - **Amortization**: Fast inference after initial training
//! - **Flexibility**: Adapts to arbitrary target distributions
//!
//! # Examples
//!
//! ```rust
//! use scirs2_core::random::neural_sampling::*;
//! use ::ndarray::Array2;
//!
//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
//! // Sample training data (small for doc test)
//! let training_data: Array2<f64> = Array2::zeros((10, 3));
//!
//! // Normalizing Flow initialization (basic example)
//! let mut flow = NormalizingFlow::new(3, 2);
//! // In real usage: flow.train(&training_data, num_epochs)?;
//!
//! // Score-based diffusion model initialization
//! let diffusion = ScoreBasedDiffusion::new(DiffusionConfig::default());
//! // In real usage: diffusion.train(&training_data)?; then diffusion.sample(...)?;
//!
//! // For this doc test, we just show initialization without expensive operations
//! println!("Neural sampling models initialized successfully");
//! # Ok(())
//! # }
//! ```

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

/// Normalizing Flow for invertible transformations
///
/// Normalizing flows learn invertible mappings between simple base distributions
/// (like Gaussian) and complex target distributions, enabling both sampling
/// and exact likelihood computation.
#[derive(Debug, Clone)]
pub struct NormalizingFlow {
    dimension: usize,
    num_layers: usize,
    flow_layers: Vec<FlowLayer>,
    base_distribution: MultivariateNormal,
    trained: bool,
    /// Per-epoch average negative log-likelihood recorded by [`Self::train`].
    training_history: Vec<f64>,
}

#[derive(Debug, Clone)]
struct FlowLayer {
    // Coupling layer parameters (simplified Real NVP-style)
    mask: Array1<bool>,
    scale_network: MLP,
    translation_network: MLP,
}

#[derive(Debug, Clone)]
struct MLP {
    // Multi-layer perceptron for flow transformations
    weights: Vec<Array2<f64>>,
    biases: Vec<Array1<f64>>,
    hidden_dims: Vec<usize>,
}

impl NormalizingFlow {
    /// Create new normalizing flow
    pub fn new(dimension: usize, num_layers: usize) -> Self {
        let mut flow_layers = Vec::new();

        for i in 0..num_layers {
            // Alternating masks for coupling layers
            let mut mask = Array1::from_elem(dimension, false);
            for j in 0..dimension {
                mask[j] = (j + i) % 2 == 0;
            }

            let hidden_dim = dimension.max(32);
            let scale_net = MLP::new(&[dimension / 2, hidden_dim, hidden_dim, dimension / 2]);
            let trans_net = MLP::new(&[dimension / 2, hidden_dim, hidden_dim, dimension / 2]);

            flow_layers.push(FlowLayer {
                mask,
                scale_network: scale_net,
                translation_network: trans_net,
            });
        }

        // Create identity covariance matrix (diagonal matrix with 1.0 on diagonal)
        let mut cov_matrix = vec![vec![0.0; dimension]; dimension];
        for i in 0..dimension {
            cov_matrix[i][i] = 1.0;
        }

        let base_distribution =
            MultivariateNormal::new(vec![0.0; dimension], cov_matrix).expect("Operation failed");

        Self {
            dimension,
            num_layers,
            flow_layers,
            base_distribution,
            trained: false,
            training_history: Vec::new(),
        }
    }

    /// Train the normalizing flow on data
    pub fn train(&mut self, training_data: &Array2<f64>, num_epochs: usize) -> Result<(), String> {
        let learning_rate = 0.001;
        let batch_size = 32;

        for epoch in 0..num_epochs {
            // Mini-batch training (simplified)
            let num_batches = training_data.nrows().div_ceil(batch_size);

            let mut epoch_loss = 0.0;
            let mut epoch_samples = 0usize;

            for batch_idx in 0..num_batches {
                let start_idx = batch_idx * batch_size;
                let end_idx = (start_idx + batch_size).min(training_data.nrows());

                let batch = training_data.slice(s![start_idx..end_idx, ..]);

                // Forward pass: compute negative log-likelihood
                let mut total_loss = 0.0;
                for i in 0..batch.nrows() {
                    let x = batch.row(i).to_owned();
                    let (z, log_det_jacobian) = self.forward(&x)?;

                    // Base distribution log probability
                    let log_prob_z = self.base_distribution.log_probability(&z.to_vec())?;
                    let log_prob_x = log_prob_z + log_det_jacobian;

                    total_loss -= log_prob_x; // Negative log-likelihood
                }

                epoch_loss += total_loss;
                epoch_samples += batch.nrows();

                // Backward pass (simplified gradient computation)
                self.update_parameters(learning_rate, &batch)?;
            }

            // Average negative log-likelihood per sample for this epoch.
            let avg_loss = if epoch_samples > 0 {
                epoch_loss / epoch_samples as f64
            } else {
                0.0
            };
            self.training_history.push(avg_loss);

            if epoch % 100 == 0 {
                println!("Epoch {epoch}: loss = {avg_loss:.6}");
            }
        }

        self.trained = true;
        Ok(())
    }

    /// Returns the per-epoch training loss history recorded by [`Self::train`].
    ///
    /// Each entry is the average negative log-likelihood over all samples
    /// processed during the corresponding epoch (one entry is pushed per
    /// epoch, in order).
    pub fn training_history(&self) -> &[f64] {
        &self.training_history
    }

    /// Forward transformation: x -> z
    fn forward(&self, x: &Array1<f64>) -> Result<(Array1<f64>, f64), String> {
        let mut z = x.clone();
        let mut log_det_jacobian = 0.0;

        for layer in &self.flow_layers {
            let (new_z, log_det) = layer.forward(&z)?;
            z = new_z;
            log_det_jacobian += log_det;
        }

        Ok((z, log_det_jacobian))
    }

    /// Inverse transformation: z -> x (for sampling)
    fn inverse(&self, z: &Array1<f64>) -> Result<Array1<f64>, String> {
        let mut x = z.clone();

        // Apply layers in reverse order
        for layer in self.flow_layers.iter().rev() {
            x = layer.inverse(&x)?;
        }

        Ok(x)
    }

    /// Sample from the learned distribution
    pub fn sample(&self, num_samples: usize, seed: u64) -> Result<Array2<f64>, String> {
        if !self.trained {
            return Err("Flow must be trained before sampling".to_string());
        }

        let mut rng = seeded_rng(seed);
        let mut samples = Array2::zeros((num_samples, self.dimension));

        for i in 0..num_samples {
            // Sample from base distribution
            let z = Array1::from_vec(self.base_distribution.sample(&mut rng));

            // Transform through flow
            let x = self.inverse(&z)?;

            for j in 0..self.dimension {
                samples[[i, j]] = x[j];
            }
        }

        Ok(samples)
    }

    /// Compute log probability of data points
    pub fn log_probability(&self, x: &Array1<f64>) -> Result<f64, String> {
        if !self.trained {
            return Err("Flow must be trained before computing probabilities".to_string());
        }

        let (z, log_det_jacobian) = self.forward(x)?;
        let log_prob_z = self.base_distribution.log_probability(&z.to_vec())?;
        Ok(log_prob_z + log_det_jacobian)
    }

    /// Update parameters (simplified gradient descent)
    fn update_parameters(
        &mut self,
        learning_rate: f64,
        batch: &crate::ndarray::ArrayView2<f64>,
    ) -> Result<(), String> {
        // Simplified parameter update - in practice would use automatic differentiation
        for layer in &mut self.flow_layers {
            layer.update_parameters(learning_rate, batch)?;
        }
        Ok(())
    }
}

impl FlowLayer {
    /// Forward pass through coupling layer
    fn forward(&self, x: &Array1<f64>) -> Result<(Array1<f64>, f64), String> {
        let mut y = x.clone();
        let mut log_det_jacobian = 0.0;

        // Split input according to mask
        let x_unchanged: Vec<f64> = x
            .iter()
            .enumerate()
            .filter(|(i, _)| self.mask[*i])
            .map(|(_, &val)| val)
            .collect();

        let x_to_transform: Vec<f64> = x
            .iter()
            .enumerate()
            .filter(|(i, _)| !self.mask[*i])
            .map(|(_, &val)| val)
            .collect();

        if !x_unchanged.is_empty() && !x_to_transform.is_empty() {
            // Compute scale and translation
            let scale = self
                .scale_network
                .forward(&Array1::from_vec(x_unchanged.clone()))?;
            let translation = self
                .translation_network
                .forward(&Array1::from_vec(x_unchanged))?;

            // Apply transformation
            let mut transform_idx = 0;
            for (i, &masked) in self.mask.iter().enumerate() {
                if !masked && transform_idx < scale.len() && transform_idx < translation.len() {
                    let s = scale[transform_idx];
                    let t = translation[transform_idx];
                    y[i] = x_to_transform[transform_idx] * s.exp() + t;
                    log_det_jacobian += s;
                    transform_idx += 1;
                }
            }
        }

        Ok((y, log_det_jacobian))
    }

    /// Inverse pass through coupling layer
    fn inverse(&self, y: &Array1<f64>) -> Result<Array1<f64>, String> {
        let mut x = y.clone();

        // Split input according to mask
        let y_unchanged: Vec<f64> = y
            .iter()
            .enumerate()
            .filter(|(i, _)| self.mask[*i])
            .map(|(_, &val)| val)
            .collect();

        if !y_unchanged.is_empty() {
            // Compute scale and translation
            let scale = self
                .scale_network
                .forward(&Array1::from_vec(y_unchanged.clone()))?;
            let translation = self
                .translation_network
                .forward(&Array1::from_vec(y_unchanged))?;

            // Apply inverse transformation
            let mut transform_idx = 0;
            for (i, &masked) in self.mask.iter().enumerate() {
                if !masked && transform_idx < scale.len() && transform_idx < translation.len() {
                    let s = scale[transform_idx];
                    let t = translation[transform_idx];
                    x[i] = (y[i] - t) * (-s).exp();
                    transform_idx += 1;
                }
            }
        }

        Ok(x)
    }

    /// Update layer parameters
    fn update_parameters(
        &mut self,
        learning_rate: f64,
        _batch: &crate::ndarray::ArrayView2<f64>,
    ) -> Result<(), String> {
        // Simplified parameter update
        self.scale_network.update_parameters(learning_rate)?;
        self.translation_network.update_parameters(learning_rate)?;
        Ok(())
    }
}

impl MLP {
    /// Create new MLP
    fn new(layer_sizes: &[usize]) -> Self {
        let mut weights = Vec::new();
        let mut biases = Vec::new();

        for i in 0..layer_sizes.len() - 1 {
            let w = Array2::zeros((layer_sizes[i + 1], layer_sizes[i]));
            let b = Array1::zeros(layer_sizes[i + 1]);
            weights.push(w);
            biases.push(b);
        }

        Self {
            weights,
            biases,
            hidden_dims: layer_sizes[1..layer_sizes.len() - 1].to_vec(),
        }
    }

    /// Forward pass through MLP
    fn forward(&self, input: &Array1<f64>) -> Result<Array1<f64>, String> {
        let mut x = input.clone();

        for (i, (weight, bias)) in self.weights.iter().zip(self.biases.iter()).enumerate() {
            // Linear transformation
            let mut output = Array1::zeros(weight.nrows());
            for j in 0..weight.nrows() {
                let mut sum = bias[j];
                for k in 0..weight.ncols() {
                    if k < x.len() {
                        sum += weight[[j, k]] * x[k];
                    }
                }
                output[j] = sum;
            }

            // Activation function (ReLU for hidden layers, linear for output)
            if i < self.weights.len() - 1 {
                for elem in output.iter_mut() {
                    *elem = elem.max(0.0); // ReLU
                }
            }

            x = output;
        }

        Ok(x)
    }

    /// Update parameters (simplified)
    fn update_parameters(&mut self, _learning_rate: f64) -> Result<(), String> {
        // Simplified parameter update - would implement proper backpropagation
        Ok(())
    }
}

/// Score-based diffusion model for high-quality sample generation
#[derive(Debug)]
pub struct ScoreBasedDiffusion {
    config: DiffusionConfig,
    score_network: ScoreNetwork,
    noise_schedule: NoiseSchedule,
    trained: bool,
    /// Per-epoch average denoising score-matching loss recorded by
    /// [`Self::train`].
    training_history: Vec<f64>,
}

#[derive(Debug, Clone)]
pub struct DiffusionConfig {
    pub dimension: usize,
    pub num_timesteps: usize,
    pub beta_start: f64,
    pub beta_end: f64,
    pub hidden_dims: Vec<usize>,
}

impl Default for DiffusionConfig {
    fn default() -> Self {
        Self {
            dimension: 10,
            num_timesteps: 1000,
            beta_start: 1e-4,
            beta_end: 0.02,
            hidden_dims: vec![128, 256, 128],
        }
    }
}

#[derive(Debug)]
struct ScoreNetwork {
    // Neural network for score function estimation
    mlp: MLP,
    time_embedding_dim: usize,
}

#[derive(Debug)]
struct NoiseSchedule {
    betas: Vec<f64>,
    alphas: Vec<f64>,
    alpha_bars: Vec<f64>,
}

impl ScoreBasedDiffusion {
    /// Create new diffusion model
    pub fn new(config: DiffusionConfig) -> Self {
        let time_embedding_dim = 64;
        let input_dim = config.dimension + time_embedding_dim;

        let mut network_dims = vec![input_dim];
        network_dims.extend_from_slice(&config.hidden_dims);
        network_dims.push(config.dimension);

        let score_network = ScoreNetwork {
            mlp: MLP::new(&network_dims),
            time_embedding_dim,
        };

        let noise_schedule =
            NoiseSchedule::new(config.num_timesteps, config.beta_start, config.beta_end);

        Self {
            config,
            score_network,
            noise_schedule,
            trained: false,
            training_history: Vec::new(),
        }
    }

    /// Train the diffusion model
    pub fn train(&mut self, training_data: &Array2<f64>) -> Result<(), String> {
        let num_epochs = 1000;
        let batch_size = 32;

        for epoch in 0..num_epochs {
            let mut epoch_loss = 0.0;
            let mut num_batches_processed = 0usize;

            // Denoising score matching training
            for _ in 0..training_data.nrows().div_ceil(batch_size) {
                // Sample random timesteps
                let mut rng = seeded_rng(42 + epoch as u64);
                let t = rng
                    .sample(Uniform::new(0, self.config.num_timesteps).expect("Operation failed"));

                // Sample noise and create noisy data
                let noise = self.sample_noise(training_data.nrows(), &mut rng)?;
                let noisy_data = self.add_noise(training_data, &noise, t)?;

                // Train score network to predict noise
                epoch_loss += self.score_network.train_step(&noisy_data, &noise, t)?;
                num_batches_processed += 1;
            }

            // Average denoising score-matching MSE loss for this epoch.
            let avg_loss = if num_batches_processed > 0 {
                epoch_loss / num_batches_processed as f64
            } else {
                0.0
            };
            self.training_history.push(avg_loss);

            if epoch % 100 == 0 {
                println!("Epoch {epoch}: loss = {avg_loss:.6}");
            }
        }

        self.trained = true;
        Ok(())
    }

    /// Returns the per-epoch training loss history recorded by [`Self::train`].
    ///
    /// Each entry is the average denoising score-matching mean-squared-error
    /// loss over all batches processed during the corresponding epoch (one
    /// entry is pushed per epoch, in order).
    pub fn training_history(&self) -> &[f64] {
        &self.training_history
    }

    /// Sample from the diffusion model using DDPM
    pub fn sample(
        &self,
        num_samples: usize,
        num_timesteps: usize,
        seed: u64,
    ) -> Result<Array2<f64>, String> {
        if !self.trained {
            return Err("Model must be trained before sampling".to_string());
        }

        let mut rng = seeded_rng(seed);
        let mut samples = Array2::zeros((num_samples, self.config.dimension));

        // Start from pure noise
        for i in 0..num_samples {
            for j in 0..self.config.dimension {
                samples[[i, j]] = rng.sample(Normal::new(0.0, 1.0).expect("Operation failed"));
            }
        }

        // Reverse diffusion process
        let timestep_stride = self.config.num_timesteps / num_timesteps;

        for t in (0..num_timesteps).rev() {
            let actual_t = t * timestep_stride;

            // Predict noise using score network
            let predicted_noise = self.score_network.predict(&samples, actual_t)?;

            // Update samples using DDPM update rule
            samples = self.ddpm_update(&samples, &predicted_noise, actual_t, &mut rng)?;
        }

        Ok(samples)
    }

    /// Sample noise
    fn sample_noise(
        &self,
        batch_size: usize,
        rng: &mut Random<rand::rngs::StdRng>,
    ) -> Result<Array2<f64>, String> {
        let mut noise = Array2::zeros((batch_size, self.config.dimension));
        for i in 0..batch_size {
            for j in 0..self.config.dimension {
                noise[[i, j]] = rng.sample(Normal::new(0.0, 1.0).expect("Operation failed"));
            }
        }
        Ok(noise)
    }

    /// Add noise according to diffusion schedule
    fn add_noise(
        &self,
        data: &Array2<f64>,
        noise: &Array2<f64>,
        t: usize,
    ) -> Result<Array2<f64>, String> {
        let alpha_bar = self.noise_schedule.alpha_bars[t];
        let mut noisy_data = Array2::zeros(data.raw_dim());

        for i in 0..data.nrows() {
            for j in 0..data.ncols() {
                noisy_data[[i, j]] =
                    alpha_bar.sqrt() * data[[i, j]] + (1.0 - alpha_bar).sqrt() * noise[[i, j]];
            }
        }

        Ok(noisy_data)
    }

    /// DDPM update step
    fn ddpm_update(
        &self,
        x_t: &Array2<f64>,
        predicted_noise: &Array2<f64>,
        t: usize,
        rng: &mut Random<rand::rngs::StdRng>,
    ) -> Result<Array2<f64>, String> {
        let alpha = self.noise_schedule.alphas[t];
        let alpha_bar = self.noise_schedule.alpha_bars[t];
        let beta = self.noise_schedule.betas[t];

        let mut x_t_minus_1 = Array2::zeros(x_t.raw_dim());

        for i in 0..x_t.nrows() {
            for j in 0..x_t.ncols() {
                // Mean of reverse process
                let mean_coeff = 1.0 / alpha.sqrt();
                let noise_coeff = beta / (1.0 - alpha_bar).sqrt();
                let mean = mean_coeff * (x_t[[i, j]] - noise_coeff * predicted_noise[[i, j]]);

                // Add noise (except for final step)
                let noise = if t > 0 {
                    rng.sample(Normal::new(0.0, beta.sqrt()).expect("Operation failed"))
                } else {
                    0.0
                };

                x_t_minus_1[[i, j]] = mean + noise;
            }
        }

        Ok(x_t_minus_1)
    }
}

impl NoiseSchedule {
    fn new(num_timesteps: usize, beta_start: f64, beta_end: f64) -> Self {
        let mut betas = Vec::with_capacity(num_timesteps);
        let mut alphas = Vec::with_capacity(num_timesteps);
        let mut alpha_bars = Vec::with_capacity(num_timesteps);

        // Linear beta schedule
        for i in 0..num_timesteps {
            let beta =
                beta_start + (beta_end - beta_start) * (i as f64) / (num_timesteps as f64 - 1.0);
            let alpha = 1.0 - beta;

            betas.push(beta);
            alphas.push(alpha);

            // Cumulative product for alpha_bar
            let alpha_bar = if i == 0 {
                alpha
            } else {
                alpha_bars[i - 1] * alpha
            };
            alpha_bars.push(alpha_bar);
        }

        Self {
            betas,
            alphas,
            alpha_bars,
        }
    }
}

impl ScoreNetwork {
    /// Train step for score network
    ///
    /// Returns the mean squared error between the network's predicted noise
    /// and the true noise added at timestep `t` — the standard denoising
    /// score-matching training objective.
    fn train_step(
        &mut self,
        noisy_data: &Array2<f64>,
        target_noise: &Array2<f64>,
        t: usize,
    ) -> Result<f64, String> {
        // Simplified training step - would implement proper backpropagation
        let mut squared_error_sum = 0.0;
        let mut count = 0usize;

        for i in 0..noisy_data.nrows() {
            let input = self.prepare_input(&noisy_data.row(i).to_owned(), t)?;
            let predicted = self.mlp.forward(&input)?;
            let target = target_noise.row(i);

            // Compute loss and update parameters
            for j in 0..predicted.len().min(target.len()) {
                let diff = predicted[j] - target[j];
                squared_error_sum += diff * diff;
                count += 1;
            }
        }

        Ok(if count > 0 {
            squared_error_sum / count as f64
        } else {
            0.0
        })
    }

    /// Predict noise at given timestep
    fn predict(&self, x: &Array2<f64>, t: usize) -> Result<Array2<f64>, String> {
        let mut predictions = Array2::zeros(x.raw_dim());

        for i in 0..x.nrows() {
            let input = self.prepare_input(&x.row(i).to_owned(), t)?;
            let pred = self.mlp.forward(&input)?;

            for j in 0..pred.len().min(x.ncols()) {
                predictions[[i, j]] = pred[j];
            }
        }

        Ok(predictions)
    }

    /// Prepare input with time embedding
    fn prepare_input(&self, x: &Array1<f64>, t: usize) -> Result<Array1<f64>, String> {
        // Simple time embedding (sinusoidal)
        let mut time_emb = Array1::zeros(self.time_embedding_dim);
        for i in 0..self.time_embedding_dim {
            let freq = 2.0 * std::f64::consts::PI * (t as f64)
                / (10000.0_f64.powf(2.0 * (i as f64) / (self.time_embedding_dim as f64)));
            time_emb[i] = if i % 2 == 0 { freq.sin() } else { freq.cos() };
        }

        // Concatenate data and time embedding
        let mut input = Array1::zeros(x.len() + time_emb.len());
        for i in 0..x.len() {
            input[i] = x[i];
        }
        for i in 0..time_emb.len() {
            input[x.len() + i] = time_emb[i];
        }

        Ok(input)
    }
}

/// Energy-Based Model for flexible probability modeling
#[derive(Debug)]
pub struct EnergyBasedModel {
    energy_network: MLP,
    dimension: usize,
    temperature: f64,
    mcmc_steps: usize,
    /// Per-epoch average contrastive-divergence loss recorded by
    /// [`Self::train`].
    training_history: Vec<f64>,
}

impl EnergyBasedModel {
    /// Create new energy-based model
    pub fn new(dimension: usize, hidden_dims: &[usize]) -> Self {
        let mut network_dims = vec![dimension];
        network_dims.extend_from_slice(hidden_dims);
        network_dims.push(1); // Single energy output

        Self {
            energy_network: MLP::new(&network_dims),
            dimension,
            temperature: 1.0,
            mcmc_steps: 100,
            training_history: Vec::new(),
        }
    }

    /// Train using contrastive divergence
    pub fn train(&mut self, training_data: &Array2<f64>, num_epochs: usize) -> Result<(), String> {
        for epoch in 0..num_epochs {
            let mut epoch_loss = 0.0;

            for i in 0..training_data.nrows() {
                let positive_sample = training_data.row(i).to_owned();

                // Generate negative sample using MCMC
                let negative_sample = self.sample_mcmc(&positive_sample, self.mcmc_steps)?;

                // Contrastive divergence update
                epoch_loss +=
                    self.contrastive_divergence_update(&positive_sample, &negative_sample)?;
            }

            // Average contrastive divergence loss (positive minus negative
            // energy) for this epoch.
            let avg_loss = if training_data.nrows() > 0 {
                epoch_loss / training_data.nrows() as f64
            } else {
                0.0
            };
            self.training_history.push(avg_loss);

            if epoch % 100 == 0 {
                println!("Epoch {epoch}: loss = {avg_loss:.6}");
            }
        }

        Ok(())
    }

    /// Returns the per-epoch training loss history recorded by [`Self::train`].
    ///
    /// Each entry is the average contrastive-divergence loss (positive-sample
    /// energy minus negative-sample energy) over all training rows processed
    /// during the corresponding epoch (one entry is pushed per epoch, in
    /// order).
    pub fn training_history(&self) -> &[f64] {
        &self.training_history
    }

    /// Sample using Langevin dynamics
    pub fn sample(
        &self,
        num_samples: usize,
        num_steps: usize,
        seed: u64,
    ) -> Result<Array2<f64>, String> {
        let mut rng = seeded_rng(seed);
        let mut samples = Array2::zeros((num_samples, self.dimension));

        for i in 0..num_samples {
            // Initialize with random noise
            let mut x = Array1::zeros(self.dimension);
            for j in 0..self.dimension {
                x[j] = rng.sample(Normal::new(0.0, 1.0).expect("Operation failed"));
            }

            // Langevin dynamics
            x = self.sample_mcmc(&x, num_steps)?;

            for j in 0..self.dimension {
                samples[[i, j]] = x[j];
            }
        }

        Ok(samples)
    }

    /// MCMC sampling using Langevin dynamics
    fn sample_mcmc(&self, initial: &Array1<f64>, num_steps: usize) -> Result<Array1<f64>, String> {
        let mut x = initial.clone();
        let step_size = 0.01;
        let mut rng = seeded_rng(42);

        for _ in 0..num_steps {
            // Compute energy gradient
            let grad = self.energy_gradient(&x)?;

            // Langevin dynamics update
            for i in 0..self.dimension {
                let noise = rng.sample(
                    Normal::new(0.0, (2.0_f64 * step_size).sqrt()).expect("Operation failed"),
                );
                x[i] -= step_size * grad[i] + noise;
            }
        }

        Ok(x)
    }

    /// Compute energy gradient (numerical differentiation)
    fn energy_gradient(&self, x: &Array1<f64>) -> Result<Array1<f64>, String> {
        let mut gradient = Array1::zeros(self.dimension);
        let epsilon = 1e-5;

        for i in 0..self.dimension {
            let mut x_plus = x.clone();
            let mut x_minus = x.clone();
            x_plus[i] += epsilon;
            x_minus[i] -= epsilon;

            let energy_plus = self.energy_network.forward(&x_plus)?[0];
            let energy_minus = self.energy_network.forward(&x_minus)?[0];

            gradient[i] = (energy_plus - energy_minus) / (2.0 * epsilon);
        }

        Ok(gradient)
    }

    /// Contrastive divergence parameter update
    ///
    /// Returns the contrastive divergence loss (positive-sample energy minus
    /// negative-sample energy) — the quantity contrastive divergence training
    /// drives down (a well-fit model assigns lower energy to real/positive
    /// samples than to MCMC-generated/negative samples).
    fn contrastive_divergence_update(
        &mut self,
        positive: &Array1<f64>,
        negative: &Array1<f64>,
    ) -> Result<f64, String> {
        // Simplified parameter update - would implement proper gradients
        let pos_energy = self.energy_network.forward(positive)?;
        let neg_energy = self.energy_network.forward(negative)?;

        // Update parameters to decrease positive energy and increase negative energy
        // (Implementation would use automatic differentiation)
        let cd_loss = pos_energy[0] - neg_energy[0];

        Ok(cd_loss)
    }
}

/// Neural Posterior Estimation for amortized Bayesian inference
#[derive(Debug)]
pub struct NeuralPosteriorEstimation {
    posterior_network: MLP,
    prior_network: MLP,
    observation_dim: usize,
    parameter_dim: usize,
    trained: bool,
    /// Per-epoch average Gaussian negative log-likelihood loss recorded by
    /// [`Self::train`].
    training_history: Vec<f64>,
}

impl NeuralPosteriorEstimation {
    /// Create new neural posterior estimator
    pub fn new(observation_dim: usize, parameter_dim: usize, hidden_dims: &[usize]) -> Self {
        // Network that takes observations and outputs posterior parameters
        let mut posterior_dims = vec![observation_dim];
        posterior_dims.extend_from_slice(hidden_dims);
        posterior_dims.push(parameter_dim * 2); // Mean and variance

        // Network that samples from prior
        let mut prior_dims = vec![parameter_dim];
        prior_dims.extend_from_slice(hidden_dims);
        prior_dims.push(parameter_dim);

        Self {
            posterior_network: MLP::new(&posterior_dims),
            prior_network: MLP::new(&prior_dims),
            observation_dim,
            parameter_dim,
            trained: false,
            training_history: Vec::new(),
        }
    }

    /// Train using simulation-based inference
    pub fn train(
        &mut self,
        simulator: impl Fn(&Array1<f64>) -> Array1<f64>,
        num_simulations: usize,
    ) -> Result<(), String> {
        let mut rng = seeded_rng(42);

        for epoch in 0..1000 {
            let mut epoch_loss = 0.0;
            let mut num_steps = 0usize;

            for _ in 0..num_simulations / 1000 {
                // Sample from prior
                let mut theta = Array1::zeros(self.parameter_dim);
                for i in 0..self.parameter_dim {
                    theta[i] = rng.sample(Normal::new(0.0, 1.0).expect("Operation failed"));
                }

                // Simulate observation
                let x = simulator(&theta);

                // Train posterior network
                epoch_loss += self.train_posterior_step(&x, &theta)?;
                num_steps += 1;
            }

            // Average Gaussian negative log-likelihood loss for this epoch.
            let avg_loss = if num_steps > 0 {
                epoch_loss / num_steps as f64
            } else {
                0.0
            };
            self.training_history.push(avg_loss);

            if epoch % 100 == 0 {
                println!("Epoch {epoch}: loss = {avg_loss:.6}");
            }
        }

        self.trained = true;
        Ok(())
    }

    /// Returns the per-epoch training loss history recorded by [`Self::train`].
    ///
    /// Each entry is the average Gaussian negative log-likelihood loss over
    /// all simulations processed during the corresponding epoch (one entry is
    /// pushed per epoch, in order — `train` always runs a fixed 1000-epoch
    /// schedule).
    pub fn training_history(&self) -> &[f64] {
        &self.training_history
    }

    /// Estimate posterior given observation
    pub fn posterior(
        &self,
        observation: &Array1<f64>,
        num_samples: usize,
        seed: u64,
    ) -> Result<Array2<f64>, String> {
        if !self.trained {
            return Err("Model must be trained before inference".to_string());
        }

        // Get posterior parameters from network
        let posterior_params = self.posterior_network.forward(observation)?;

        let mean_start = 0;
        let var_start = self.parameter_dim;

        let mut rng = seeded_rng(seed);
        let mut samples = Array2::zeros((num_samples, self.parameter_dim));

        for i in 0..num_samples {
            for j in 0..self.parameter_dim {
                let mean = posterior_params[mean_start + j];
                let var = posterior_params[var_start + j].exp(); // Ensure positive variance

                samples[[i, j]] =
                    rng.sample(Normal::new(mean, var.sqrt()).expect("Operation failed"));
            }
        }

        Ok(samples)
    }

    /// Train posterior network step
    ///
    /// Returns the Gaussian negative log-likelihood of `true_parameter` under
    /// the predicted posterior mean/variance — the amortized inference loss
    /// this step is meant to minimize.
    fn train_posterior_step(
        &mut self,
        observation: &Array1<f64>,
        true_parameter: &Array1<f64>,
    ) -> Result<f64, String> {
        // Get predicted posterior parameters
        let predicted_params = self.posterior_network.forward(observation)?;

        let mean_start = 0;
        let var_start = self.parameter_dim;

        // Compute loss (negative log-likelihood) and update
        // (Implementation would use automatic differentiation)
        let mut nll = 0.0;
        for j in 0..self.parameter_dim {
            let mean = predicted_params[mean_start + j];
            let var = predicted_params[var_start + j].exp().max(1e-12); // Ensure positive variance
            let diff = true_parameter[j] - mean;
            nll += 0.5 * diff * diff / var + 0.5 * var.ln();
        }
        nll += 0.5 * (self.parameter_dim as f64) * (2.0 * std::f64::consts::PI).ln();

        Ok(nll)
    }
}

// Helper trait for extending base distribution with log probability
trait LogProbability {
    fn log_probability(&self, x: &[f64]) -> Result<f64, String>;
}

impl LogProbability for MultivariateNormal {
    fn log_probability(&self, x: &[f64]) -> Result<f64, String> {
        if x.len() != self.dimension() {
            return Err("Dimension mismatch".to_string());
        }

        // Simplified log probability computation
        let mut log_prob = 0.0;
        for &xi in x {
            log_prob += -0.5 * xi * xi; // Assume standard normal for simplicity
        }
        log_prob += -0.5 * (x.len() as f64) * (2.0 * std::f64::consts::PI).ln();

        Ok(log_prob)
    }
}

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

    #[test]
    fn test_normalizing_flow_creation() {
        let flow = NormalizingFlow::new(5, 3);
        assert_eq!(flow.dimension, 5);
        assert_eq!(flow.num_layers, 3);
        assert!(!flow.trained);
    }

    #[test]
    fn test_diffusion_model_creation() {
        let config = DiffusionConfig {
            dimension: 10,
            num_timesteps: 100,
            beta_start: 1e-4,
            beta_end: 0.02,
            hidden_dims: vec![32, 64, 32],
        };

        let diffusion = ScoreBasedDiffusion::new(config);
        assert_eq!(diffusion.config.dimension, 10);
        assert_eq!(diffusion.config.num_timesteps, 100);
    }

    #[test]
    fn test_energy_based_model() {
        let ebm = EnergyBasedModel::new(5, &[32, 32]);
        assert_eq!(ebm.dimension, 5);
        assert_eq!(ebm.mcmc_steps, 100);
    }

    #[test]
    fn test_neural_posterior_estimation() {
        let npe = NeuralPosteriorEstimation::new(10, 5, &[32, 32]);
        assert_eq!(npe.observation_dim, 10);
        assert_eq!(npe.parameter_dim, 5);
        assert!(!npe.trained);
    }

    #[test]
    fn test_mlp_forward() {
        let mlp = MLP::new(&[3, 5, 2]);
        let input = Array1::from_vec(vec![1.0, 2.0, 3.0]);
        let output = mlp.forward(&input).expect("Operation failed");
        assert_eq!(output.len(), 2);
    }

    #[test]
    fn test_noise_schedule() {
        let schedule = NoiseSchedule::new(100, 1e-4, 0.02);
        assert_eq!(schedule.betas.len(), 100);
        assert_eq!(schedule.alphas.len(), 100);
        assert_eq!(schedule.alpha_bars.len(), 100);

        // Check that alpha_bars are decreasing
        for i in 1..schedule.alpha_bars.len() {
            assert!(schedule.alpha_bars[i] <= schedule.alpha_bars[i - 1]);
        }
    }

    #[test]
    fn test_normalizing_flow_training_history() {
        let mut flow = NormalizingFlow::new(4, 2);
        let training_data =
            Array2::from_shape_fn((8, 4), |(i, j)| (i as f64 * 0.3 + j as f64 * 0.1) - 1.0);
        let num_epochs = 5;

        flow.train(&training_data, num_epochs)
            .expect("training should succeed on well-formed synthetic data");

        let history = flow.training_history();
        assert_eq!(
            history.len(),
            num_epochs,
            "training_history should have exactly one entry per epoch"
        );
        for (i, &loss) in history.iter().enumerate() {
            assert!(
                loss.is_finite(),
                "epoch {i} loss should be finite, got {loss}"
            );
        }
    }

    #[test]
    fn test_score_based_diffusion_training_history() {
        let config = DiffusionConfig {
            dimension: 3,
            num_timesteps: 20,
            beta_start: 1e-4,
            beta_end: 0.02,
            hidden_dims: vec![8, 8],
        };
        let mut diffusion = ScoreBasedDiffusion::new(config);
        let training_data = Array2::from_shape_fn((6, 3), |(i, j)| (i as f64 - j as f64) * 0.2);

        diffusion
            .train(&training_data)
            .expect("training should succeed on well-formed synthetic data");

        let history = diffusion.training_history();
        // `train` has no epoch parameter and always runs a fixed 1000-epoch
        // schedule internally.
        assert_eq!(
            history.len(),
            1000,
            "training_history should have exactly one entry per (fixed) epoch"
        );
        for (i, &loss) in history.iter().enumerate() {
            assert!(
                loss.is_finite(),
                "epoch {i} loss should be finite, got {loss}"
            );
        }
    }

    #[test]
    fn test_energy_based_model_training_history() {
        let mut ebm = EnergyBasedModel::new(3, &[8, 8]);
        let training_data = Array2::from_shape_fn((4, 3), |(i, j)| (i as f64 - j as f64) * 0.15);
        let num_epochs = 4;

        ebm.train(&training_data, num_epochs)
            .expect("training should succeed on well-formed synthetic data");

        let history = ebm.training_history();
        assert_eq!(
            history.len(),
            num_epochs,
            "training_history should have exactly one entry per epoch"
        );
        for (i, &loss) in history.iter().enumerate() {
            assert!(
                loss.is_finite(),
                "epoch {i} loss should be finite, got {loss}"
            );
        }
    }

    #[test]
    fn test_neural_posterior_estimation_training_history() {
        let mut npe = NeuralPosteriorEstimation::new(4, 2, &[8, 8]);
        let simulator = |theta: &Array1<f64>| -> Array1<f64> {
            Array1::from_vec(vec![
                theta[0],
                theta[1],
                theta[0] + theta[1],
                theta[0] - theta[1],
            ])
        };

        npe.train(simulator, 2000)
            .expect("training should succeed on well-formed synthetic simulator");

        let history = npe.training_history();
        // `train` has no epoch parameter and always runs a fixed 1000-epoch
        // schedule internally.
        assert_eq!(
            history.len(),
            1000,
            "training_history should have exactly one entry per (fixed) epoch"
        );
        for (i, &loss) in history.iter().enumerate() {
            assert!(
                loss.is_finite(),
                "epoch {i} loss should be finite, got {loss}"
            );
        }
    }
}