//! LCM Sampler (Latent Consistency Models)
//!
//! Ultra-fast 1-4 step sampling through consistency model formulation.
//! Supports consistency distillation from pre-trained diffusion models and
//! classifier-free guidance for high-quality generation.

use crate::error::{DiffusionError, Result};
use crate::noise::NoiseGenerator;
use rtx_tensor::Tensor;

/// LCM prediction types for different parameterizations
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum LCMPredictionType {
    /// Predict noise (epsilon parameterization)
    Epsilon,
    /// Predict velocity (v-parameterization)
    VPrediction,
    /// Predict data directly (x0 parameterization)
    Data,
}

/// Configuration for LCM sampler
#[derive(Debug, Clone)]
pub struct LCMConfig {
    /// Number of sampling steps (1-4 for LCM)
    pub num_steps: u32,
    /// Classifier-free guidance scale
    pub guidance_scale: f32,
    /// Consistency loss weight
    pub consistency_weight: f32,
    /// Prediction type
    pub prediction_type: LCMPredictionType,
    /// Solver order (1-3)
    pub solver_order: u8,
    /// Training mode flag
    pub training_mode: bool,
    /// Use Karras sigmas
    pub use_karras_sigmas: bool,
    /// Distillation loss weight
    pub distillation_weight: f32,
}

/// Statistics for LCM sampler
#[derive(Debug, Default)]
pub struct LCMStats {
    /// Number of function evaluations
    pub nfe: usize,
    /// Consistency loss sum
    pub consistency_loss_sum: f32,
    /// Distillation loss sum
    pub distillation_loss_sum: f32,
    /// Average sampling time (ms)
    pub avg_sampling_time_ms: f32,
    /// Number of sampling runs
    pub num_runs: usize,
}

/// LCM Sampler for ultra-fast 1-4 step generation
pub struct LCMSampler {
    config: LCMConfig,
    noise_generator: NoiseGenerator,
    stats: LCMStats,
    /// Cached sigmas for efficiency
    cached_sigmas: Option<Vec<f32>>,
}

impl Default for LCMConfig {
    fn default() -> Self {
        Self {
            num_steps: 4,
            guidance_scale: 7.5,
            consistency_weight: 1.0,
            prediction_type: LCMPredictionType::Epsilon,
            solver_order: 2,
            training_mode: false,
            use_karras_sigmas: false,
            distillation_weight: 0.5,
        }
    }
}

impl LCMSampler {
    /// Create a new LCM sampler
    pub fn new(config: LCMConfig, noise_generator: NoiseGenerator) -> Result<Self> {
        // Validation
        if config.num_steps == 0 || config.num_steps > 4 {
            return Err(DiffusionError::Scheduler {
                message: "LCM num_steps must be between 1 and 4".to_string(),
            });
        }

        if config.guidance_scale < 0.0 {
            return Err(DiffusionError::Scheduler {
                message: "Guidance scale must be non-negative".to_string(),
            });
        }

        Ok(Self {
            config,
            noise_generator,
            stats: LCMStats::default(),
            cached_sigmas: None,
        })
    }

    /// Get current statistics
    pub fn stats(&self) -> &LCMStats {
        &self.stats
    }

    /// Sample using consistency model (single/multi-step)
    pub fn sample<F>(&mut self, x_t: &Tensor, timestep: u32, model_fn: F) -> Result<Tensor>
    where
        F: Fn(&Tensor, u32) -> Result<Tensor>,
    {
        let start_time = std::time::Instant::now();
        let sigmas = self.get_sigmas(timestep)?;
        let mut x = x_t.clone();

        // Multi-step LCM sampling with consistency formulation
        for i in 0..self.config.num_steps {
            let current_timestep = self.sigma_to_timestep(sigmas[i as usize])?;
            let model_output = model_fn(&x, current_timestep)?;
            self.stats.nfe += 1;

            let x0_pred = self.convert_prediction(&x, &model_output, current_timestep)?;

            x = if i < self.config.num_steps - 1 {
                let next_sigma = sigmas[(i + 1) as usize];
                self.consistency_step(&x0_pred, &x, sigmas[i as usize], next_sigma)?
            } else {
                x0_pred
            };
        }

        self.update_timing_stats(start_time.elapsed());
        Ok(x)
    }

    /// Sample with classifier-free guidance
    pub fn sample_guided<F>(
        &mut self,
        x_t: &Tensor,
        timestep: u32,
        guided_model_fn: F,
    ) -> Result<Tensor>
    where
        F: Fn(&Tensor, u32) -> Result<(Tensor, Tensor)>,
    {
        let start_time = std::time::Instant::now();
        let sigmas = self.get_sigmas(timestep)?;
        let mut x = x_t.clone();

        for i in 0..self.config.num_steps {
            let current_timestep = self.sigma_to_timestep(sigmas[i as usize])?;
            let (uncond_output, cond_output) = guided_model_fn(&x, current_timestep)?;
            self.stats.nfe += 2;

            let guided_output = self.apply_cfg(&uncond_output, &cond_output)?;
            let x0_pred = self.convert_prediction(&x, &guided_output, current_timestep)?;

            x = if i < self.config.num_steps - 1 {
                let next_sigma = sigmas[(i + 1) as usize];
                self.consistency_step(&x0_pred, &x, sigmas[i as usize], next_sigma)?
            } else {
                x0_pred
            };
        }

        self.update_timing_stats(start_time.elapsed());
        Ok(x)
    }

    /// Compute consistency loss for training
    pub fn compute_consistency_loss(
        &self,
        x0: &Tensor,
        noise: &Tensor,
        t1: u32,
        t2: u32,
    ) -> Result<Tensor> {
        let x_t1 = self.noise_generator.add_noise(x0, noise, t1)?;
        let x_t2 = self.noise_generator.add_noise(x0, noise, t2)?;

        let f_t1 = self.consistency_function(&x_t1, t1)?;
        let f_t2 = self.consistency_function(&x_t2, t2)?;

        let diff = f_t1.subtract(&f_t2)?;
        let loss = diff.pow_scalar(2.0)?.mean(&[], false)?;
        loss.scalar_mul(self.config.consistency_weight)
            .map_err(DiffusionError::Tensor)
    }

    /// Compute distillation loss from teacher model
    pub fn compute_distillation_loss(
        &self,
        x0: &Tensor,
        teacher_output: &Tensor,
        student_output: &Tensor,
        timestep: u32,
    ) -> Result<Tensor> {
        let teacher_x0 = self.convert_prediction(x0, teacher_output, timestep)?;
        let student_x0 = self.convert_prediction(x0, student_output, timestep)?;

        let diff = teacher_x0.subtract(&student_x0)?;
        let loss = diff.pow_scalar(2.0)?.mean(&[], false)?;
        loss.scalar_mul(self.config.distillation_weight)
            .map_err(DiffusionError::Tensor)
    }

    /// Convert model prediction to x0 based on prediction type
    pub fn convert_prediction(
        &self,
        x_t: &Tensor,
        model_output: &Tensor,
        timestep: u32,
    ) -> Result<Tensor> {
        match self.config.prediction_type {
            LCMPredictionType::Epsilon => {
                // x0 = (x_t - sqrt(1-alpha_cumprod) * epsilon) / sqrt(alpha_cumprod)
                let (sqrt_alpha_cumprod, sqrt_one_minus_alpha_cumprod, _, _) =
                    self.noise_generator.get_schedule_params(timestep)?;

                let scaled_noise = model_output.scalar_mul(sqrt_one_minus_alpha_cumprod)?;
                let x_minus_noise = x_t.subtract(&scaled_noise)?;
                x_minus_noise
                    .scalar_mul(1.0 / sqrt_alpha_cumprod)
                    .map_err(DiffusionError::Tensor)
            }
            LCMPredictionType::VPrediction => {
                // v-parameterization: x0 = sqrt(alpha_cumprod) * x_t - sqrt(1-alpha_cumprod) * v
                let (sqrt_alpha_cumprod, sqrt_one_minus_alpha_cumprod, _, _) =
                    self.noise_generator.get_schedule_params(timestep)?;

                let x_component = x_t.scalar_mul(sqrt_alpha_cumprod)?;
                let v_component = model_output.scalar_mul(sqrt_one_minus_alpha_cumprod)?;
                x_component
                    .subtract(&v_component)
                    .map_err(DiffusionError::Tensor)
            }
            LCMPredictionType::Data => {
                // Direct x0 prediction
                Ok(model_output.clone())
            }
        }
    }

    /// Get Karras sigmas for improved sampling
    pub fn get_karras_sigmas(&self, num_steps: u32) -> Result<Vec<f32>> {
        if let Some(ref cached) = self.cached_sigmas {
            if cached.len() == (num_steps + 1) as usize {
                return Ok(cached.clone());
            }
        }

        let sigma_min: f32 = 0.002;
        let sigma_max: f32 = 80.0;
        let rho: f32 = 7.0;

        let mut sigmas = Vec::with_capacity((num_steps + 1) as usize);

        for i in 0..=num_steps {
            let u = i as f32 / num_steps as f32;
            let sigma: f32 = (sigma_max.powf(1.0_f32 / rho)
                + u * (sigma_min.powf(1.0_f32 / rho) - sigma_max.powf(1.0_f32 / rho)))
            .powf(rho);
            sigmas.push(sigma);
        }

        Ok(sigmas)
    }

    // Helper methods (optimized implementations)

    fn get_sigmas(&self, timestep: u32) -> Result<Vec<f32>> {
        if self.config.use_karras_sigmas {
            self.get_karras_sigmas(self.config.num_steps)
        } else {
            self.get_linear_sigmas(self.config.num_steps, timestep)
        }
    }

    fn get_linear_sigmas(&self, num_steps: u32, max_timestep: u32) -> Result<Vec<f32>> {
        let mut sigmas = Vec::with_capacity((num_steps + 1) as usize);

        for i in 0..=num_steps {
            let t = max_timestep as f32 * (1.0 - i as f32 / num_steps as f32);
            let (_, sqrt_one_minus_alpha_cumprod, _, _) =
                self.noise_generator.get_schedule_params(t as u32)?;
            sigmas.push(sqrt_one_minus_alpha_cumprod);
        }

        Ok(sigmas)
    }

    fn update_timing_stats(&mut self, elapsed: std::time::Duration) {
        self.stats.num_runs += 1;
        let total_time = self.stats.avg_sampling_time_ms * (self.stats.num_runs - 1) as f32
            + elapsed.as_millis() as f32;
        self.stats.avg_sampling_time_ms = total_time / self.stats.num_runs as f32;
    }

    fn sigma_to_timestep(&self, sigma: f32) -> Result<u32> {
        // Simplified mapping - in practice would be more sophisticated
        let timestep = (sigma * 1000.0) as u32;
        Ok(timestep.min(self.noise_generator.num_timesteps() - 1))
    }

    fn consistency_function(&self, x_t: &Tensor, timestep: u32) -> Result<Tensor> {
        // Simplified consistency function: f(x_t, t) = x_t / (1 + sigma(t))
        let (_, sqrt_one_minus_alpha_cumprod, _, _) =
            self.noise_generator.get_schedule_params(timestep)?;
        let scale = 1.0 / (1.0 + sqrt_one_minus_alpha_cumprod);
        x_t.scalar_mul(scale).map_err(DiffusionError::Tensor)
    }

    fn consistency_step(
        &self,
        x0_pred: &Tensor,
        x_t: &Tensor,
        sigma_curr: f32,
        sigma_next: f32,
    ) -> Result<Tensor> {
        // Simplified consistency step: blend between prediction and current state
        let alpha = sigma_next / sigma_curr;
        let x0_component = x0_pred.scalar_mul(1.0 - alpha)?;
        let xt_component = x_t.scalar_mul(alpha)?;
        x0_component
            .add(&xt_component)
            .map_err(DiffusionError::Tensor)
    }

    fn apply_cfg(&self, uncond_output: &Tensor, cond_output: &Tensor) -> Result<Tensor> {
        // CFG: output = uncond + guidance_scale * (cond - uncond)
        let diff = cond_output.subtract(uncond_output)?;
        let scaled_diff = diff.scalar_mul(self.config.guidance_scale)?;
        uncond_output
            .add(&scaled_diff)
            .map_err(DiffusionError::Tensor)
    }

    /// Advanced consistency training loss with boundary conditions
    pub fn compute_advanced_consistency_loss(
        &self,
        x0: &Tensor,
        noise: &Tensor,
        timesteps: &[u32],
    ) -> Result<Tensor> {
        if timesteps.len() < 2 {
            return Err(DiffusionError::Scheduler {
                message: "Need at least 2 timesteps for consistency loss".to_string(),
            });
        }

        let mut total_loss = None;
        for window in timesteps.windows(2) {
            let t1 = window[0];
            let t2 = window[1];
            let window_loss = self.compute_consistency_loss(x0, noise, t1, t2)?;

            total_loss = match total_loss {
                None => Some(window_loss),
                Some(acc) => Some(acc.add(&window_loss)?),
            };
        }

        total_loss
            .unwrap()
            .scalar_mul(1.0 / (timesteps.len() - 1) as f32)
            .map_err(DiffusionError::Tensor)
    }

    /// Adaptive sampling with dynamic step adjustment
    pub fn sample_adaptive<F>(
        &mut self,
        x_t: &Tensor,
        timestep: u32,
        model_fn: F,
        error_threshold: f32,
    ) -> Result<Tensor>
    where
        F: Fn(&Tensor, u32) -> Result<Tensor>,
    {
        let mut current_steps = self.config.num_steps.max(1);
        let mut x = x_t.clone();

        while current_steps <= 4 {
            let backup_steps = self.config.num_steps;
            self.config.num_steps = current_steps;

            let result = self.sample(x_t, timestep, &model_fn);
            self.config.num_steps = backup_steps;

            match result {
                Ok(sampled) => {
                    // Simple error estimation based on change magnitude
                    let change = sampled.subtract(&x)?.pow_scalar(2.0)?.mean(&[], false)?;
                    let error = change.item()?;

                    if error < error_threshold || current_steps == 4 {
                        return Ok(sampled);
                    }
                    x = sampled;
                    current_steps += 1;
                }
                Err(e) => return Err(e),
            }
        }

        Ok(x)
    }

    /// Get solver-specific sigma schedule
    pub fn get_solver_sigmas(&self, timestep: u32) -> Result<Vec<f32>> {
        match self.config.solver_order {
            1 => self.get_euler_sigmas(timestep),
            2 => self.get_heun_sigmas(timestep),
            3 => self.get_dpm_sigmas(timestep),
            _ => self.get_sigmas(timestep),
        }
    }

    fn get_euler_sigmas(&self, timestep: u32) -> Result<Vec<f32>> {
        // Euler method sigmas (linear spacing)
        self.get_linear_sigmas(self.config.num_steps, timestep)
    }

    fn get_heun_sigmas(&self, timestep: u32) -> Result<Vec<f32>> {
        // Heun method sigmas (better for 2-step)
        if self.config.use_karras_sigmas {
            self.get_karras_sigmas(self.config.num_steps)
        } else {
            self.get_linear_sigmas(self.config.num_steps, timestep)
        }
    }

    fn get_dpm_sigmas(&self, timestep: u32) -> Result<Vec<f32>> {
        // DPM-style sigma spacing for higher order
        let mut sigmas = self.get_karras_sigmas(self.config.num_steps)?;
        // Apply DPM weighting
        for sigma in &mut sigmas {
            *sigma = sigma.sqrt(); // Square root weighting for better stability
        }
        Ok(sigmas)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::noise::NoiseSchedule;
    use rtx_tensor::Device;

    #[test]
    fn test_lcm_config_creation() {
        let config = LCMConfig {
            num_steps: 4,
            guidance_scale: 7.5,
            consistency_weight: 1.0,
            prediction_type: LCMPredictionType::Epsilon,
            solver_order: 2,
            training_mode: false,
            use_karras_sigmas: false,
            distillation_weight: 0.5,
        };

        assert_eq!(config.num_steps, 4);
        assert_eq!(config.guidance_scale, 7.5);
        assert_eq!(config.consistency_weight, 1.0);
    }

    #[test]
    fn test_lcm_sampler_creation() {
        let noise_gen = NoiseGenerator::new(
            NoiseSchedule::Linear {
                beta_start: 0.0001,
                beta_end: 0.02,
            },
            1000,
            Some(42),
        )
        .unwrap();

        let config = LCMConfig::default();
        let sampler = LCMSampler::new(config, noise_gen);

        assert!(sampler.is_ok());
        let sampler = sampler.unwrap();
        assert_eq!(sampler.stats().nfe, 0);
    }

    #[test]
    fn test_consistency_function_single_step() {
        let noise_gen =
            NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();

        let config = LCMConfig {
            num_steps: 1,
            ..Default::default()
        };
        let mut sampler = LCMSampler::new(config, noise_gen).unwrap();

        // Mock model output function
        let model_fn = |x: &Tensor, t: u32| -> Result<Tensor> {
            // Simple mock: return scaled input
            let scale = (t as f32 / 1000.0) * 0.1;
            Ok(x.scalar_mul(scale)?)
        };

        let input_shape = vec![1, 3, 32, 32];
        let device = Device::cuda(0).unwrap_or(Device::default());
        let latent = Tensor::randn(&input_shape, &device).unwrap();

        let result = sampler.sample(&latent, 999, model_fn);
        assert!(result.is_ok());

        let denoised = result.unwrap();
        assert_eq!(denoised.shape(), input_shape);
        assert_eq!(sampler.stats().nfe, 1);
    }

    #[test]
    fn test_consistency_function_multi_step() {
        let noise_gen = NoiseGenerator::new(
            NoiseSchedule::Linear {
                beta_start: 0.0001,
                beta_end: 0.02,
            },
            1000,
            Some(42),
        )
        .unwrap();

        let config = LCMConfig {
            num_steps: 4,
            ..Default::default()
        };
        let mut sampler = LCMSampler::new(config, noise_gen).unwrap();

        let model_fn = |x: &Tensor, t: u32| -> Result<Tensor> {
            let scale = (t as f32 / 1000.0) * 0.1;
            Ok(x.scalar_mul(scale)?)
        };

        let input_shape = vec![2, 4, 64, 64];
        let device = Device::cuda(0).unwrap_or(Device::default());
        let latent = Tensor::randn(&input_shape, &device).unwrap();

        let result = sampler.sample(&latent, 800, model_fn);
        assert!(result.is_ok());

        let denoised = result.unwrap();
        assert_eq!(denoised.shape(), input_shape);
        assert_eq!(sampler.stats().nfe, 4);
    }

    #[test]
    fn test_classifier_free_guidance() {
        let noise_gen =
            NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();

        let config = LCMConfig {
            num_steps: 2,
            guidance_scale: 7.5,
            ..Default::default()
        };
        let mut sampler = LCMSampler::new(config, noise_gen).unwrap();

        let guided_model_fn = |x: &Tensor, t: u32| -> Result<(Tensor, Tensor)> {
            let scale = (t as f32 / 1000.0) * 0.1;
            let uncond = x.scalar_mul(scale)?;
            let cond = x.scalar_mul(scale * 1.2)?;
            Ok((uncond, cond))
        };

        let input_shape = vec![1, 4, 32, 32];
        let device = Device::cuda(0).unwrap_or(Device::default());
        let latent = Tensor::randn(&input_shape, &device).unwrap();

        let result = sampler.sample_guided(&latent, 900, guided_model_fn);
        assert!(result.is_ok());

        let denoised = result.unwrap();
        assert_eq!(denoised.shape(), input_shape);
        // Should have called model twice per step (uncond + cond)
        assert_eq!(sampler.stats().nfe, 4);
    }

    #[test]
    #[ignore = "Pre-existing Metal shader compilation issue"]
    fn test_consistency_loss_computation() {
        let noise_gen = NoiseGenerator::new(
            NoiseSchedule::Linear {
                beta_start: 0.0001,
                beta_end: 0.02,
            },
            1000,
            Some(42),
        )
        .unwrap();

        let config = LCMConfig {
            training_mode: true,
            consistency_weight: 2.0,
            ..Default::default()
        };
        let sampler = LCMSampler::new(config, noise_gen).unwrap();

        let x0_shape = vec![2, 3, 64, 64];
        let device = Device::cuda(0).unwrap_or(Device::default());
        let x0 = Tensor::randn(&x0_shape, &device).unwrap();
        let noise = Tensor::randn(&x0_shape, &device).unwrap();

        let t1 = 100;
        let t2 = 200;

        let loss = sampler.compute_consistency_loss(&x0, &noise, t1, t2);
        assert!(loss.is_ok());

        let loss_value = loss.unwrap();
        // Loss should be a scalar
        assert_eq!(loss_value.shape(), vec![]);
        assert!(loss_value.item().unwrap() >= 0.0);
    }

    #[test]
    #[ignore = "Pre-existing Metal shader compilation issue"]
    fn test_distillation_loss_computation() {
        let noise_gen =
            NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();

        let config = LCMConfig {
            training_mode: true,
            distillation_weight: 1.5,
            ..Default::default()
        };
        let sampler = LCMSampler::new(config, noise_gen).unwrap();

        let x0_shape = vec![1, 4, 32, 32];
        let device = Device::cuda(0).unwrap_or(Device::default());
        let x0 = Tensor::randn(&x0_shape, &device).unwrap();
        let teacher_output = Tensor::randn(&x0_shape, &device).unwrap();
        let student_output = Tensor::randn(&x0_shape, &device).unwrap();

        let timestep = 500;

        let loss =
            sampler.compute_distillation_loss(&x0, &teacher_output, &student_output, timestep);
        assert!(loss.is_ok());

        let loss_value = loss.unwrap();
        assert_eq!(loss_value.shape(), vec![]);
        assert!(loss_value.item().unwrap() >= 0.0);
    }

    #[test]
    fn test_prediction_type_conversion() {
        let noise_gen = NoiseGenerator::new(
            NoiseSchedule::Linear {
                beta_start: 0.0001,
                beta_end: 0.02,
            },
            1000,
            Some(42),
        )
        .unwrap();

        let config = LCMConfig {
            prediction_type: LCMPredictionType::VPrediction,
            ..Default::default()
        };
        let sampler = LCMSampler::new(config, noise_gen).unwrap();

        let x_shape = vec![1, 3, 32, 32];
        let device = Device::cuda(0).unwrap_or(Device::default());
        let x_t = Tensor::randn(&x_shape, &device).unwrap();
        let model_output = Tensor::randn(&x_shape, &device).unwrap();
        let timestep = 500;

        let x0_pred = sampler.convert_prediction(&x_t, &model_output, timestep);
        assert!(x0_pred.is_ok());
        assert_eq!(x0_pred.unwrap().shape(), x_shape);
    }

    #[test]
    fn test_karras_sigma_schedule() {
        let noise_gen =
            NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();

        let config = LCMConfig {
            use_karras_sigmas: true,
            num_steps: 4,
            ..Default::default()
        };
        let sampler = LCMSampler::new(config, noise_gen).unwrap();

        let sigmas = sampler.get_karras_sigmas(4);
        assert!(sigmas.is_ok());

        let sigma_vec = sigmas.unwrap();
        assert_eq!(sigma_vec.len(), 5); // num_steps + 1

        // Sigmas should be decreasing
        for i in 1..sigma_vec.len() {
            assert!(sigma_vec[i - 1] >= sigma_vec[i]);
        }

        // Last sigma should be close to 0
        assert!(sigma_vec.last().unwrap() < &0.01);
    }

    #[test]
    fn test_lcm_stats_tracking() {
        let noise_gen = NoiseGenerator::new(
            NoiseSchedule::Linear {
                beta_start: 0.0001,
                beta_end: 0.02,
            },
            1000,
            Some(42),
        )
        .unwrap();

        let config = LCMConfig::default();
        let mut sampler = LCMSampler::new(config, noise_gen).unwrap();

        // Initial stats
        let initial_stats = sampler.stats();
        assert_eq!(initial_stats.nfe, 0);
        assert_eq!(initial_stats.consistency_loss_sum, 0.0);
        assert_eq!(initial_stats.distillation_loss_sum, 0.0);

        // Perform sampling to update stats
        let model_fn = |x: &Tensor, _t: u32| -> Result<Tensor> { Ok(x.scalar_mul(0.1)?) };

        let device = Device::cuda(0).unwrap_or(Device::default());
        let latent = Tensor::randn(&vec![1, 3, 32, 32], &device).unwrap();
        let _ = sampler.sample(&latent, 800, model_fn);

        // Stats should be updated
        let final_stats = sampler.stats();
        assert!(final_stats.nfe > 0);
    }

    #[test]
    fn test_invalid_configurations() {
        let noise_gen = NoiseGenerator::new(
            NoiseSchedule::Linear {
                beta_start: 0.0001,
                beta_end: 0.02,
            },
            1000,
            Some(42),
        )
        .unwrap();

        // Test invalid num_steps
        let invalid_config = LCMConfig {
            num_steps: 0,
            ..Default::default()
        };
        let result = LCMSampler::new(invalid_config, noise_gen.clone());
        assert!(result.is_err());

        // Test invalid guidance_scale
        let invalid_config = LCMConfig {
            guidance_scale: -1.0,
            ..Default::default()
        };
        let result = LCMSampler::new(invalid_config, noise_gen);
        assert!(result.is_err());
    }

    #[test]
    #[ignore = "Pre-existing Metal shader compilation issue"]
    fn test_advanced_consistency_loss() {
        let noise_gen =
            NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();

        let config = LCMConfig {
            training_mode: true,
            consistency_weight: 1.5,
            ..Default::default()
        };
        let sampler = LCMSampler::new(config, noise_gen).unwrap();

        let x0_shape = vec![1, 3, 32, 32];
        let device = Device::cuda(0).unwrap_or(Device::default());
        let x0 = Tensor::randn(&x0_shape, &device).unwrap();
        let noise = Tensor::randn(&x0_shape, &device).unwrap();

        let timesteps = vec![100, 200, 300, 400];

        let loss = sampler.compute_advanced_consistency_loss(&x0, &noise, &timesteps);
        assert!(loss.is_ok());

        let loss_value = loss.unwrap();
        assert_eq!(loss_value.shape(), vec![]);
        assert!(loss_value.item().unwrap() >= 0.0);
    }

    #[test]
    #[ignore = "Pre-existing Metal shader compilation issue"]
    fn test_adaptive_sampling() {
        let noise_gen = NoiseGenerator::new(
            NoiseSchedule::Linear {
                beta_start: 0.0001,
                beta_end: 0.02,
            },
            1000,
            Some(42),
        )
        .unwrap();

        let config = LCMConfig {
            num_steps: 1,
            ..Default::default()
        };
        let mut sampler = LCMSampler::new(config, noise_gen).unwrap();

        let model_fn = |x: &Tensor, _t: u32| -> Result<Tensor> { Ok(x.scalar_mul(0.1)?) };

        let input_shape = vec![1, 3, 32, 32];
        let device = Device::cuda(0).unwrap_or(Device::default());
        let latent = Tensor::randn(&input_shape, &device).unwrap();

        let result = sampler.sample_adaptive(&latent, 800, model_fn, 0.01);
        assert!(result.is_ok());

        let denoised = result.unwrap();
        assert_eq!(denoised.shape(), input_shape);
    }

    #[test]
    fn test_solver_specific_sigmas() {
        let noise_gen =
            NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();

        // Test different solver orders
        for order in 1..=3 {
            let config = LCMConfig {
                solver_order: order,
                num_steps: 4,
                ..Default::default()
            };
            let sampler = LCMSampler::new(config, noise_gen.clone()).unwrap();

            let sigmas = sampler.get_solver_sigmas(800);
            assert!(sigmas.is_ok());

            let sigma_vec = sigmas.unwrap();
            assert_eq!(sigma_vec.len(), 5);
        }
    }
}