//! UniPC Sampler (Unified Predictor-Corrector)
//!
//! A unified predictor-corrector framework for diffusion models with fast 5-10 step sampling.
//! Supports adaptive orders, variance reduction, and efficient model evaluation caching.

use crate::error::{DiffusionError, Result};
use crate::noise::NoiseGenerator;
use rtx_tensor::Tensor;
use std::collections::VecDeque;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

#[derive(Debug, Clone)]
pub struct UniPCConfig {
    pub predictor_order: u8,
    pub corrector_order: u8,
    pub adaptive_order: bool,
    pub use_corrector: bool,
    pub prediction_type: PredictionType,
    pub variance_reduction: bool,
    pub corrector_iterations: u8,
    pub max_order: u8,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum PredictionType {
    Epsilon,
    VPrediction,
    Data,
}

#[derive(Debug, Default)]
pub struct UniPCStats {
    pub nfe: usize,
    pub predictor_steps: usize,
    pub corrector_steps: usize,
    pub cache_hits: usize,
    pub cache_misses: usize,
    pub order_adjustments: usize,
}

#[derive(Debug, Clone)]
struct CacheEntry {
    timestep: u32,
    sample_hash: u64,
    output: Tensor,
}

pub struct UniPCSampler {
    config: UniPCConfig,
    noise_generator: NoiseGenerator,
    predictor_outputs: VecDeque<Tensor>,
    corrector_outputs: VecDeque<Tensor>,
    timestep_history: VecDeque<u32>,
    sample_history: VecDeque<Tensor>,
    current_predictor_order: u8,
    current_corrector_order: u8,
    eval_cache: Vec<CacheEntry>,
    stats: UniPCStats,
}

impl UniPCSampler {
    pub fn new(config: UniPCConfig, noise_generator: NoiseGenerator) -> Result<Self> {
        if config.predictor_order == 0 || config.predictor_order > 3 {
            return Err(DiffusionError::Scheduler {
                message: "Predictor order must be 1, 2, or 3".to_string(),
            });
        }
        if config.corrector_order == 0 || config.corrector_order > 3 {
            return Err(DiffusionError::Scheduler {
                message: "Corrector order must be 1, 2, or 3".to_string(),
            });
        }
        if config.max_order == 0 || config.max_order > 3 {
            return Err(DiffusionError::Scheduler {
                message: "Max order must be 1, 2, or 3".to_string(),
            });
        }

        let current_predictor_order = if config.adaptive_order {
            1
        } else {
            config.predictor_order
        };
        let current_corrector_order = if config.adaptive_order {
            1
        } else {
            config.corrector_order
        };

        Ok(Self {
            config,
            noise_generator,
            predictor_outputs: VecDeque::new(),
            corrector_outputs: VecDeque::new(),
            timestep_history: VecDeque::new(),
            sample_history: VecDeque::new(),
            current_predictor_order,
            current_corrector_order,
            eval_cache: Vec::new(),
            stats: UniPCStats::default(),
        })
    }

    pub fn step(
        &mut self,
        model_output: &Tensor,
        timestep: u32,
        sample: &Tensor,
    ) -> Result<Tensor> {
        self.stats.nfe += 1;
        let converted_output = self.convert_prediction_type(
            model_output,
            timestep,
            sample,
            self.config.prediction_type,
            self.config.prediction_type,
        )?;
        self.update_history(&converted_output, timestep, sample)?;

        let predictor_result =
            if self.predictor_outputs.len() >= self.current_predictor_order as usize {
                let outputs: Vec<&Tensor> = self
                    .predictor_outputs
                    .iter()
                    .take(self.current_predictor_order as usize)
                    .collect();
                let timesteps: Vec<u32> = self
                    .timestep_history
                    .iter()
                    .take(self.current_predictor_order as usize)
                    .cloned()
                    .collect();
                self.predictor_step(&outputs, &timesteps, sample, self.current_predictor_order)?
            } else {
                self.predictor_step(&[&converted_output], &[timestep], sample, 1)?
            };
        self.stats.predictor_steps += 1;

        if self.config.use_corrector {
            self.stats.corrector_steps += 1;
            self.corrector_step(
                &predictor_result,
                &converted_output,
                timestep,
                self.current_corrector_order,
            )
        } else {
            Ok(predictor_result)
        }
    }

    /// Perform predictor step using Taylor expansion
    pub fn predictor_step(
        &self,
        model_outputs: &[&Tensor],
        timesteps: &[u32],
        sample: &Tensor,
        order: u8,
    ) -> Result<Tensor> {
        if model_outputs.is_empty() || timesteps.is_empty() {
            return Err(DiffusionError::Scheduler {
                message: "Empty model outputs or timesteps".to_string(),
            });
        }

        match order {
            1 => self.predictor_order_1(model_outputs[0], timesteps[0], sample),
            2 => {
                if model_outputs.len() >= 2 && timesteps.len() >= 2 {
                    self.predictor_order_2(model_outputs, timesteps, sample)
                } else {
                    self.predictor_order_1(model_outputs[0], timesteps[0], sample)
                }
            }
            3 => {
                if model_outputs.len() >= 3 && timesteps.len() >= 3 {
                    self.predictor_order_3(model_outputs, timesteps, sample)
                } else if model_outputs.len() >= 2 {
                    self.predictor_order_2(model_outputs, timesteps, sample)
                } else {
                    self.predictor_order_1(model_outputs[0], timesteps[0], sample)
                }
            }
            _ => Err(DiffusionError::Scheduler {
                message: format!("Invalid predictor order: {}", order),
            }),
        }
    }

    /// Perform corrector step with variance reduction
    pub fn corrector_step(
        &mut self,
        predicted_sample: &Tensor,
        model_output: &Tensor,
        timestep: u32,
        order: u8,
    ) -> Result<Tensor> {
        // Apply variance reduction if enabled
        let variance_reduced = if self.config.variance_reduction {
            self.apply_variance_reduction(predicted_sample, model_output, timestep)?
        } else {
            predicted_sample.clone()
        };

        // Apply corrector based on order
        match order {
            1 => {
                // Simple corrector: weighted average
                let weight = 0.95;
                let corrected = variance_reduced.scalar_mul(weight)?;
                let noise_component = model_output.scalar_mul(1.0 - weight)?;
                corrected
                    .add(&noise_component)
                    .map_err(|e| DiffusionError::TensorError(e.to_string()))
            }
            2 | 3 => {
                // Higher order corrector with more iterations
                let mut result = variance_reduced;
                for _ in 0..self.config.corrector_iterations {
                    let weight = 0.9;
                    let temp = result.scalar_mul(weight)?;
                    let correction = model_output.scalar_mul(1.0 - weight)?;
                    result = temp
                        .add(&correction)
                        .map_err(|e| DiffusionError::TensorError(e.to_string()))?;
                }
                Ok(result)
            }
            _ => Err(DiffusionError::Scheduler {
                message: format!("Invalid corrector order: {}", order),
            }),
        }
    }

    /// Configure timesteps for fast sampling (5-10 steps)
    pub fn configure_fast_timesteps(&self, num_steps: u32) -> Result<Vec<u32>> {
        let total_timesteps = self.noise_generator.num_timesteps();

        if num_steps == 0 {
            return Err(DiffusionError::Scheduler {
                message: "Number of steps must be greater than 0".to_string(),
            });
        }

        if num_steps > total_timesteps {
            return Err(DiffusionError::Scheduler {
                message: format!(
                    "Number of steps ({}) cannot exceed total timesteps ({})",
                    num_steps, total_timesteps
                ),
            });
        }

        let mut timesteps = Vec::with_capacity(num_steps as usize);

        // Use uniform spacing for simplicity, could be improved with non-uniform spacing
        let step_size = total_timesteps / num_steps;

        for i in 0..num_steps {
            let timestep = total_timesteps - 1 - (i * step_size);
            timesteps.push(timestep);
        }

        Ok(timesteps)
    }

    /// Convert between prediction types
    pub fn convert_prediction_type(
        &self,
        model_output: &Tensor,
        timestep: u32,
        sample: &Tensor,
        from: PredictionType,
        to: PredictionType,
    ) -> Result<Tensor> {
        if from == to {
            return Ok(model_output.clone());
        }

        let (sqrt_alpha_cumprod, sqrt_one_minus_alpha_cumprod, alpha_cumprod, _) =
            self.noise_generator.get_schedule_params(timestep)?;

        match (from, to) {
            (PredictionType::Epsilon, PredictionType::Data) => {
                // x0 = (x_t - sqrt(1-alpha_cumprod) * noise) / sqrt(alpha_cumprod)
                let scaled_noise = model_output
                    .scalar_mul(sqrt_one_minus_alpha_cumprod)
                    .map_err(|e| DiffusionError::TensorError(e.to_string()))?;
                let x_minus_noise = sample
                    .subtract(&scaled_noise)
                    .map_err(|e| DiffusionError::TensorError(e.to_string()))?;
                x_minus_noise
                    .scalar_mul(1.0 / sqrt_alpha_cumprod)
                    .map_err(|e| DiffusionError::TensorError(e.to_string()))
            }
            (PredictionType::Epsilon, PredictionType::VPrediction) => {
                // v = sqrt(alpha_cumprod) * noise - sqrt(1-alpha_cumprod) * x0
                // First convert to x0, then to v
                let x0 = self.convert_prediction_type(
                    model_output,
                    timestep,
                    sample,
                    from,
                    PredictionType::Data,
                )?;
                let scaled_noise = model_output
                    .scalar_mul(sqrt_alpha_cumprod)
                    .map_err(|e| DiffusionError::TensorError(e.to_string()))?;
                let scaled_x0 = x0
                    .scalar_mul(sqrt_one_minus_alpha_cumprod)
                    .map_err(|e| DiffusionError::TensorError(e.to_string()))?;
                scaled_noise
                    .subtract(&scaled_x0)
                    .map_err(|e| DiffusionError::TensorError(e.to_string()))
            }
            (PredictionType::Data, PredictionType::Epsilon) => {
                // noise = (x_t - sqrt(alpha_cumprod) * x0) / sqrt(1-alpha_cumprod)
                let scaled_x0 = model_output
                    .scalar_mul(sqrt_alpha_cumprod)
                    .map_err(|e| DiffusionError::TensorError(e.to_string()))?;
                let x_minus_x0 = sample
                    .subtract(&scaled_x0)
                    .map_err(|e| DiffusionError::TensorError(e.to_string()))?;
                x_minus_x0
                    .scalar_mul(1.0 / sqrt_one_minus_alpha_cumprod)
                    .map_err(|e| DiffusionError::TensorError(e.to_string()))
            }
            (PredictionType::VPrediction, PredictionType::Epsilon) => {
                // v = sqrt(alpha_cumprod) * noise - sqrt(1-alpha_cumprod) * x0
                // noise = (v + sqrt(1-alpha_cumprod) * x0) / sqrt(alpha_cumprod)
                // Need to solve for x0 first: x_t = sqrt(alpha_cumprod)*x0 + sqrt(1-alpha_cumprod)*noise
                // x0 = (x_t - sqrt(1-alpha_cumprod)*noise) / sqrt(alpha_cumprod)
                // Using v-prediction: x_t = sqrt(alpha_cumprod)*x0 + sqrt(1-alpha_cumprod)*((v + sqrt(1-alpha_cumprod)*x0)/sqrt(alpha_cumprod))
                let alpha_cumprod_sqrt = sqrt_alpha_cumprod;
                let one_minus_alpha_sqrt = sqrt_one_minus_alpha_cumprod;

                // Simplified approach: assume linear relationship for v-prediction conversion
                let scale = alpha_cumprod_sqrt / (alpha_cumprod_sqrt + one_minus_alpha_sqrt);
                model_output
                    .scalar_mul(scale)
                    .map_err(|e| DiffusionError::TensorError(e.to_string()))
            }
            _ => {
                // Other conversions (Data->VPrediction, VPrediction->Data) would require similar logic
                Ok(model_output.clone()) // Placeholder for now
            }
        }
    }

    /// Update history buffers for multistep methods
    pub fn update_history(
        &mut self,
        model_output: &Tensor,
        timestep: u32,
        sample: &Tensor,
    ) -> Result<()> {
        // Add to front of history buffers
        self.predictor_outputs.push_front(model_output.clone());
        self.timestep_history.push_front(timestep);
        self.sample_history.push_front(sample.clone());

        // Maintain maximum history size based on max_order
        let max_history = self.config.max_order as usize;

        while self.predictor_outputs.len() > max_history {
            self.predictor_outputs.pop_back();
        }
        while self.timestep_history.len() > max_history {
            self.timestep_history.pop_back();
        }
        while self.sample_history.len() > max_history {
            self.sample_history.pop_back();
        }

        Ok(())
    }

    /// Get cached model evaluation if available
    pub fn get_cached_output(&mut self, timestep: u32, sample: &Tensor) -> Option<&Tensor> {
        // Simple hash-based cache lookup
        let sample_hash = self.compute_tensor_hash(sample);

        for entry in &self.eval_cache {
            if entry.timestep == timestep && entry.sample_hash == sample_hash {
                self.stats.cache_hits += 1;
                return Some(&entry.output);
            }
        }

        self.stats.cache_misses += 1;
        None
    }

    /// Cache model evaluation result
    pub fn cache_output(&mut self, timestep: u32, sample: &Tensor, output: &Tensor) -> Result<()> {
        let sample_hash = self.compute_tensor_hash(sample);

        let entry = CacheEntry {
            timestep,
            sample_hash,
            output: output.clone(),
        };

        // Simple cache management: limit size to prevent unbounded growth
        const MAX_CACHE_SIZE: usize = 100;
        if self.eval_cache.len() >= MAX_CACHE_SIZE {
            self.eval_cache.remove(0); // Remove oldest entry
        }

        self.eval_cache.push(entry);
        Ok(())
    }

    /// Adaptive order selection based on error estimate
    pub fn select_adaptive_order(
        &mut self,
        error_estimate: f32,
        step_budget: u32,
    ) -> Result<(u8, u8)> {
        // Simple adaptive strategy based on error and remaining steps
        let predictor_order = if error_estimate > 0.1 || step_budget < 5 {
            1 // Use lower order for high error or few remaining steps
        } else if error_estimate > 0.05 || step_budget < 10 {
            2
        } else {
            3 // Use highest order for low error and sufficient steps
        }
        .min(self.config.max_order);

        let corrector_order = if error_estimate > 0.1 {
            1 // Simple corrector for high error
        } else {
            2 // More sophisticated corrector for lower error  
        }
        .min(self.config.max_order);

        // Update current orders if different
        if self.current_predictor_order != predictor_order
            || self.current_corrector_order != corrector_order
        {
            self.stats.order_adjustments += 1;
            self.current_predictor_order = predictor_order;
            self.current_corrector_order = corrector_order;
        }

        Ok((predictor_order, corrector_order))
    }

    /// Apply variance reduction techniques
    pub fn apply_variance_reduction(
        &self,
        predicted_sample: &Tensor,
        model_output: &Tensor,
        timestep: u32,
    ) -> Result<Tensor> {
        // Simple variance reduction using exponential moving average
        let variance = self.noise_generator.get_variance(timestep).unwrap_or(1.0);
        let reduction_factor = (1.0 / (1.0 + variance)).min(0.95);

        let reduced_sample = predicted_sample
            .scalar_mul(reduction_factor)
            .map_err(|e| DiffusionError::TensorError(e.to_string()))?;
        let noise_component = model_output
            .scalar_mul(1.0 - reduction_factor)
            .map_err(|e| DiffusionError::TensorError(e.to_string()))?;

        reduced_sample
            .add(&noise_component)
            .map_err(|e| DiffusionError::TensorError(e.to_string()))
    }

    /// Get sampling statistics
    pub fn stats(&self) -> &UniPCStats {
        &self.stats
    }

    /// Reset sampler state
    pub fn reset(&mut self) {
        self.predictor_outputs.clear();
        self.corrector_outputs.clear();
        self.timestep_history.clear();
        self.sample_history.clear();
        self.eval_cache.clear();
        self.current_predictor_order = if self.config.adaptive_order {
            1
        } else {
            self.config.predictor_order
        };
        self.current_corrector_order = if self.config.adaptive_order {
            1
        } else {
            self.config.corrector_order
        };
        self.stats = UniPCStats::default();
    }

    // Helper methods for predictor orders
    fn predictor_order_1(
        &self,
        model_output: &Tensor,
        timestep: u32,
        sample: &Tensor,
    ) -> Result<Tensor> {
        // First order predictor (essentially Euler method)
        let (sqrt_alpha_cumprod, sqrt_one_minus_alpha_cumprod, alpha_cumprod, alpha_cumprod_prev) =
            self.noise_generator.get_schedule_params(timestep)?;

        // Simple first-order step
        let scale = (alpha_cumprod_prev / alpha_cumprod).sqrt();
        let scaled_sample = sample
            .scalar_mul(scale)
            .map_err(|e| DiffusionError::TensorError(e.to_string()))?;

        // Add model prediction contribution
        let pred_scale = (1.0 - alpha_cumprod_prev).sqrt() - scale * sqrt_one_minus_alpha_cumprod;
        let pred_component = model_output
            .scalar_mul(pred_scale)
            .map_err(|e| DiffusionError::TensorError(e.to_string()))?;

        scaled_sample
            .add(&pred_component)
            .map_err(|e| DiffusionError::TensorError(e.to_string()))
    }

    fn predictor_order_2(
        &self,
        model_outputs: &[&Tensor],
        timesteps: &[u32],
        sample: &Tensor,
    ) -> Result<Tensor> {
        // Second order predictor using linear combination
        let result = self.predictor_order_1(model_outputs[0], timesteps[0], sample)?;

        // Add second-order correction
        if model_outputs.len() >= 2 {
            let correction_weight = 0.5;
            let correction = model_outputs[1]
                .scalar_mul(correction_weight)
                .map_err(|e| DiffusionError::TensorError(e.to_string()))?;
            result
                .add(&correction)
                .map_err(|e| DiffusionError::TensorError(e.to_string()))
        } else {
            Ok(result)
        }
    }

    fn predictor_order_3(
        &self,
        model_outputs: &[&Tensor],
        timesteps: &[u32],
        sample: &Tensor,
    ) -> Result<Tensor> {
        // Third order predictor using quadratic combination
        let result = self.predictor_order_2(model_outputs, timesteps, sample)?;

        // Add third-order correction
        if model_outputs.len() >= 3 {
            let correction_weight = 0.25;
            let correction = model_outputs[2]
                .scalar_mul(correction_weight)
                .map_err(|e| DiffusionError::TensorError(e.to_string()))?;
            result
                .add(&correction)
                .map_err(|e| DiffusionError::TensorError(e.to_string()))
        } else {
            Ok(result)
        }
    }

    fn compute_tensor_hash(&self, tensor: &Tensor) -> u64 {
        // Simple hash based on tensor shape - in practice would use actual data
        let mut hasher = DefaultHasher::new();
        tensor.shape().hash(&mut hasher);
        hasher.finish()
    }
}

impl Default for UniPCConfig {
    fn default() -> Self {
        Self {
            predictor_order: 2,
            corrector_order: 1,
            adaptive_order: false,
            use_corrector: true,
            prediction_type: PredictionType::Epsilon,
            variance_reduction: true,
            corrector_iterations: 1,
            max_order: 3,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::noise::{NoiseGenerator, NoiseSchedule};

    fn create_test_noise_generator() -> NoiseGenerator {
        NoiseGenerator::new(
            NoiseSchedule::Linear {
                beta_start: 0.0001,
                beta_end: 0.02,
            },
            1000,
            Some(42),
        )
        .unwrap()
    }
    fn create_test_tensor(shape: Vec<usize>) -> Tensor {
        let total_size = shape.iter().product::<usize>();
        let data: Vec<f32> = (0..total_size).map(|i| i as f32 * 0.01).collect();
        Tensor::new(data, shape).unwrap()
    }

    #[test]
    fn test_unipc_creation_and_validation() {
        let noise_gen = create_test_noise_generator();
        let sampler = UniPCSampler::new(UniPCConfig::default(), noise_gen.clone()).unwrap();
        assert_eq!(sampler.config.predictor_order, 2);
        assert_eq!(sampler.config.corrector_order, 1);
        assert!(sampler.config.use_corrector);

        let config = UniPCConfig {
            predictor_order: 3,
            corrector_order: 2,
            adaptive_order: true,
            use_corrector: true,
            prediction_type: PredictionType::VPrediction,
            variance_reduction: true,
            corrector_iterations: 2,
            max_order: 3,
        };
        let sampler = UniPCSampler::new(config, noise_gen).unwrap();
        assert_eq!(sampler.config.predictor_order, 3);
        assert!(sampler.config.adaptive_order);
    }

    #[test]
    fn test_unipc_config_validation() {
        let noise_gen = create_test_noise_generator();
        assert!(
            UniPCSampler::new(
                UniPCConfig {
                    predictor_order: 0,
                    ..Default::default()
                },
                noise_gen.clone()
            )
            .is_err()
        );
        assert!(
            UniPCSampler::new(
                UniPCConfig {
                    corrector_order: 4,
                    ..Default::default()
                },
                noise_gen.clone()
            )
            .is_err()
        );
        assert!(
            UniPCSampler::new(
                UniPCConfig {
                    max_order: 0,
                    ..Default::default()
                },
                noise_gen
            )
            .is_err()
        );
    }

    #[test]
    fn test_predictor_orders() {
        let mut sampler =
            UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
        let sample = create_test_tensor(vec![1, 3, 32, 32]);
        let model_output = create_test_tensor(vec![1, 3, 32, 32]);

        for (order, outputs, timesteps) in [
            (1, vec![&model_output], vec![500]),
            (2, vec![&model_output, &model_output], vec![500, 400]),
            (
                3,
                vec![&model_output, &model_output, &model_output],
                vec![500, 400, 300],
            ),
        ] {
            let result = sampler.predictor_step(&outputs, &timesteps, &sample, order);
            assert!(result.is_ok());
            assert_eq!(result.unwrap().shape(), sample.shape());
        }
    }

    #[test]
    fn test_corrector_functionality() {
        let mut sampler = UniPCSampler::new(
            UniPCConfig {
                use_corrector: true,
                corrector_order: 2,
                corrector_iterations: 2,
                ..Default::default()
            },
            create_test_noise_generator(),
        )
        .unwrap();
        let sample = create_test_tensor(vec![1, 3, 32, 32]);
        let model_output = create_test_tensor(vec![1, 3, 32, 32]);

        for order in [1, 2, 3] {
            let result = sampler.corrector_step(&sample, &model_output, 500, order);
            assert!(result.is_ok());
            assert_eq!(result.unwrap().shape(), sample.shape());
        }
    }

    #[test]
    fn test_full_sampling_step() {
        let mut sampler =
            UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
        let sample = create_test_tensor(vec![1, 3, 32, 32]);
        let model_output = create_test_tensor(vec![1, 3, 32, 32]);

        let result = sampler.step(&model_output, 500, &sample).unwrap();
        assert_eq!(result.shape(), sample.shape());
        assert_eq!(sampler.stats().nfe, 1);
        assert!(sampler.stats().predictor_steps > 0);

        let mut current_sample = sample;
        let timesteps = vec![500u32, 400, 300, 200, 100];
        for timestep in timesteps {
            current_sample = sampler
                .step(&model_output, timestep, &current_sample)
                .unwrap();
        }
        assert!(sampler.stats().nfe >= 5);
    }

    #[test]
    fn test_prediction_type_conversion() {
        let sampler =
            UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
        let sample = create_test_tensor(vec![1, 3, 32, 32]);
        let model_output = create_test_tensor(vec![1, 3, 32, 32]);
        let timestep = 500;

        let conversions = [
            (PredictionType::Epsilon, PredictionType::VPrediction),
            (PredictionType::Epsilon, PredictionType::Data),
            (PredictionType::Epsilon, PredictionType::Epsilon),
        ];
        for (from, to) in conversions {
            let result =
                sampler.convert_prediction_type(&model_output, timestep, &sample, from, to);
            assert!(result.is_ok());
            assert_eq!(result.unwrap().shape(), model_output.shape());
        }
    }

    #[test]
    fn test_cache_functionality() {
        let mut sampler =
            UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
        let sample = create_test_tensor(vec![1, 3, 32, 32]);
        let model_output = create_test_tensor(vec![1, 3, 32, 32]);

        assert!(sampler.get_cached_output(500, &sample).is_none());
        assert_eq!(sampler.stats().cache_misses, 1);
        assert!(sampler.cache_output(500, &sample, &model_output).is_ok());

        // Get cached output and verify shape before checking stats
        let cached_shape = {
            let cached = sampler.get_cached_output(500, &sample);
            assert!(cached.is_some());
            cached.unwrap().shape().to_vec()
        };

        assert_eq!(sampler.stats().cache_hits, 1);
        assert_eq!(cached_shape, model_output.shape().dims());
    }

    #[test]
    fn test_fast_timesteps_configuration() {
        let sampler =
            UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
        for num_steps in [5, 7, 10, 15] {
            let ts = sampler.configure_fast_timesteps(num_steps).unwrap();
            assert_eq!(ts.len(), num_steps as usize);
            for i in 1..ts.len() {
                assert!(ts[i - 1] > ts[i]);
            }
        }
        assert!(sampler.configure_fast_timesteps(0).is_err());
        assert!(sampler.configure_fast_timesteps(2000).is_err());
    }

    #[test]
    fn test_adaptive_order_selection() {
        let mut sampler = UniPCSampler::new(
            UniPCConfig {
                adaptive_order: true,
                max_order: 3,
                ..Default::default()
            },
            create_test_noise_generator(),
        )
        .unwrap();

        for (error, steps) in [(0.01, 10), (0.1, 5)] {
            let (pred_order, corr_order) = sampler.select_adaptive_order(error, steps).unwrap();
            assert!(pred_order >= 1 && pred_order <= 3 && corr_order >= 1 && corr_order <= 3);
        }
        let (pred_low, _) = sampler.select_adaptive_order(0.5, 5).unwrap();
        let (pred_high, _) = sampler.select_adaptive_order(0.01, 10).unwrap();
        assert!(pred_low <= pred_high);
    }

    #[test]
    fn test_variance_reduction() {
        let sampler = UniPCSampler::new(
            UniPCConfig {
                variance_reduction: true,
                ..Default::default()
            },
            create_test_noise_generator(),
        )
        .unwrap();
        let sample = create_test_tensor(vec![1, 3, 32, 32]);
        let model_output = create_test_tensor(vec![1, 3, 32, 32]);
        let reduced = sampler
            .apply_variance_reduction(&sample, &model_output, 500)
            .unwrap();
        assert_eq!(reduced.shape(), sample.shape());
    }

    #[test]
    fn test_history_management() {
        let mut sampler = UniPCSampler::new(
            UniPCConfig {
                predictor_order: 3,
                ..Default::default()
            },
            create_test_noise_generator(),
        )
        .unwrap();
        let sample = create_test_tensor(vec![1, 3, 32, 32]);
        let model_output = create_test_tensor(vec![1, 3, 32, 32]);

        for timestep in (100..=500).step_by(50) {
            assert!(
                sampler
                    .update_history(&model_output, timestep, &sample)
                    .is_ok()
            );
        }
        assert!(sampler.timestep_history.len() <= sampler.config.max_order as usize);
        assert!(sampler.predictor_outputs.len() <= sampler.config.max_order as usize);
    }

    #[test]
    fn test_reset_functionality() {
        let mut sampler =
            UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
        let sample = create_test_tensor(vec![1, 3, 32, 32]);
        let model_output = create_test_tensor(vec![1, 3, 32, 32]);

        let _result = sampler.step(&model_output, 500, &sample).unwrap();
        let _cache = sampler.cache_output(500, &sample, &model_output);
        assert!(sampler.stats().nfe > 0 && !sampler.eval_cache.is_empty());

        sampler.reset();
        assert_eq!(sampler.stats().nfe, 0);
        assert!(sampler.eval_cache.is_empty() && sampler.timestep_history.is_empty());
    }
}