//! Quantization Module for Model Compression and Acceleration
//!
//! This module implements various quantization schemes (INT8, INT4, FP8) with
//! calibration methods to maintain model accuracy while reducing memory usage
//! and improving inference performance.
//!
//! Key features:
//! - Multiple quantization schemes: INT8, INT4, FP8 E4M3, FP8 E5M2
//! - Dynamic quantization with calibration
//! - Per-tensor and per-channel granularity
//! - Outlier handling and accuracy validation
//! - Mixed precision support for different layers

use crate::{InferenceError, InferenceResult};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::{debug, info};

/// Data types supported by quantization
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DataType {
    F32,
    F16,
    BF16,
    INT8,
    INT4,
    FP8E4M3,
    FP8E5M2,
}

/// Quantization schemes available
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QuantizationScheme {
    INT8,
    INT4,
    FP8E4M3,
    FP8E5M2,
}

/// Scaling methods for quantization
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScalingMethod {
    Symmetric,
    Asymmetric,
}

/// Calibration methods for dynamic quantization
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CalibrationMethod {
    MinMax,
    Percentile,
    KLDivergence,
    MSE,
}

/// Quantization granularity
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QuantizationGranularity {
    PerTensor,
    PerChannel,
    PerGroup,
}

/// Group size for grouped quantization
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GroupSize {
    G32,
    G64,
    G128,
    G256,
}

impl GroupSize {
    #[must_use]
    pub fn size(&self) -> usize {
        match self {
            Self::G32 => 32,
            Self::G64 => 64,
            Self::G128 => 128,
            Self::G256 => 256,
        }
    }
}

/// Outlier handling strategies
#[derive(Debug, Clone, PartialEq)]
pub enum OutlierHandling {
    None,
    Clipping { percentile: f32 },
    Separate { threshold: f32 },
}

/// Optimization levels for quantization performance
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OptimizationLevel {
    None,
    Basic,
    Aggressive,
}

/// Configuration for quantization operations
#[derive(Debug, Clone)]
pub struct QuantizationConfig {
    pub scheme: QuantizationScheme,
    pub scaling_method: ScalingMethod,
    pub calibration_method: CalibrationMethod,
    pub calibration_samples: usize,
    pub accuracy_threshold: f32,
    pub granularity: QuantizationGranularity,
    pub group_size: Option<GroupSize>,
    pub dynamic_range_optimization: bool,
    pub outlier_handling: OutlierHandling,
    pub optimization_level: OptimizationLevel,
}

impl QuantizationConfig {
    /// Create a new configuration with defaults
    #[must_use]
    pub fn new() -> Self {
        Self {
            scheme: QuantizationScheme::INT8,
            scaling_method: ScalingMethod::Symmetric,
            calibration_method: CalibrationMethod::MinMax,
            calibration_samples: 512,
            accuracy_threshold: 0.01,
            granularity: QuantizationGranularity::PerTensor,
            group_size: None,
            dynamic_range_optimization: false,
            outlier_handling: OutlierHandling::None,
            optimization_level: OptimizationLevel::Basic,
        }
    }

    /// Set quantization scheme
    #[must_use]
    pub fn with_scheme(mut self, scheme: QuantizationScheme) -> Self {
        self.scheme = scheme;
        self
    }

    /// Set scaling method
    #[must_use]
    pub fn with_scaling_method(mut self, method: ScalingMethod) -> Self {
        self.scaling_method = method;
        self
    }

    /// Set calibration method
    #[must_use]
    pub fn with_calibration_method(mut self, method: CalibrationMethod) -> Self {
        self.calibration_method = method;
        self
    }

    /// Set number of calibration samples
    #[must_use]
    pub fn with_calibration_samples(mut self, samples: usize) -> Self {
        self.calibration_samples = samples;
        self
    }

    /// Set accuracy threshold for validation
    #[must_use]
    pub fn with_accuracy_threshold(mut self, threshold: f32) -> Self {
        self.accuracy_threshold = threshold;
        self
    }

    /// Set quantization granularity
    #[must_use]
    pub fn with_granularity(mut self, granularity: QuantizationGranularity) -> Self {
        self.granularity = granularity;
        self
    }

    /// Set group size for grouped quantization
    #[must_use]
    pub fn with_group_size(mut self, size: GroupSize) -> Self {
        self.group_size = Some(size);
        self
    }

    /// Enable dynamic range optimization
    #[must_use]
    pub fn with_dynamic_range_optimization(mut self, enabled: bool) -> Self {
        self.dynamic_range_optimization = enabled;
        self
    }

    /// Set outlier handling method
    #[must_use]
    pub fn with_outlier_handling(mut self, handling: OutlierHandling) -> Self {
        self.outlier_handling = handling;
        self
    }

    /// Set optimization level
    #[must_use]
    pub fn with_optimization_level(mut self, level: OptimizationLevel) -> Self {
        self.optimization_level = level;
        self
    }

    /// Validate configuration parameters
    pub fn validate(&self) -> InferenceResult<()> {
        if self.calibration_samples == 0 {
            return Err(InferenceError::invalid_request(
                "calibration_samples must be greater than 0",
            ));
        }

        if !(0.0..=1.0).contains(&self.accuracy_threshold) {
            return Err(InferenceError::invalid_request(
                "accuracy_threshold must be between 0.0 and 1.0",
            ));
        }

        if matches!(self.granularity, QuantizationGranularity::PerGroup)
            && self.group_size.is_none()
        {
            return Err(InferenceError::invalid_request(
                "group_size must be specified for per-group quantization",
            ));
        }

        if let OutlierHandling::Clipping { percentile } = self.outlier_handling
            && !(0.0..=100.0).contains(&percentile)
        {
            return Err(InferenceError::invalid_request(
                "outlier clipping percentile must be between 0.0 and 100.0",
            ));
        }

        Ok(())
    }
}

impl Default for QuantizationConfig {
    fn default() -> Self {
        Self::new()
    }
}

/// Test tensor representation for quantization operations
#[derive(Debug, Clone)]
pub struct TestTensor {
    pub data: Vec<f32>,
    pub shape: Vec<usize>,
    pub dtype: DataType,
}

impl TestTensor {
    #[must_use]
    pub fn new(data: Vec<f32>, shape: Vec<usize>) -> Self {
        Self {
            data,
            shape,
            dtype: DataType::F32,
        }
    }

    #[must_use]
    pub fn with_dtype(mut self, dtype: DataType) -> Self {
        self.dtype = dtype;
        self
    }

    #[must_use]
    pub fn num_elements(&self) -> usize {
        self.shape.iter().product()
    }

    #[must_use]
    pub fn from_random(shape: Vec<usize>, min: f32, max: f32) -> Self {
        let num_elements = shape.iter().product();
        let mut data = Vec::with_capacity(num_elements);
        for _ in 0..num_elements {
            data.push(min + fastrand::f32() * (max - min));
        }
        Self::new(data, shape)
    }

    #[must_use]
    pub fn storage_size_bytes(&self) -> usize {
        match self.dtype {
            DataType::F32 => self.data.len() * 4,
            DataType::F16 => self.data.len() * 2,
            DataType::BF16 => self.data.len() * 2,
            DataType::INT8 => self.data.len(),
            DataType::INT4 => self.data.len().div_ceil(2), // 4 bits per element
            DataType::FP8E4M3 => self.data.len(),
            DataType::FP8E5M2 => self.data.len(),
        }
    }

    fn validate(&self) -> InferenceResult<()> {
        if self.data.is_empty() {
            return Err(InferenceError::invalid_request("Tensor data is empty"));
        }

        let expected_elements = self.shape.iter().product::<usize>();
        if self.data.len() != expected_elements {
            return Err(InferenceError::invalid_request(format!(
                "Data length ({}) doesn't match shape ({:?})",
                self.data.len(),
                self.shape
            )));
        }

        // Check for invalid values
        for (i, &value) in self.data.iter().enumerate() {
            if !value.is_finite() {
                return Err(InferenceError::invalid_request(format!(
                    "Non-finite value {value} at index {i}"
                )));
            }
        }

        Ok(())
    }
}

/// Quantization parameters computed during calibration
#[derive(Debug, Clone)]
pub struct QuantizationParams {
    pub scale: f32,
    pub zero_point: i32,
    pub min_value: f32,
    pub max_value: f32,
    pub scheme: QuantizationScheme,
}

impl QuantizationParams {
    #[must_use]
    pub fn new(scheme: QuantizationScheme) -> Self {
        Self {
            scale: 1.0,
            zero_point: 0,
            min_value: 0.0,
            max_value: 0.0,
            scheme,
        }
    }
}

/// Main quantization implementation
pub struct Quantizer {
    config: QuantizationConfig,
    params: Option<QuantizationParams>,
}

impl Quantizer {
    /// Create a new quantizer
    #[must_use]
    pub fn new(config: QuantizationConfig) -> Self {
        Self {
            config,
            params: None,
        }
    }

    /// Quantize a tensor
    pub async fn quantize(&self, tensor: &TestTensor) -> InferenceResult<TestTensor> {
        tensor.validate()?;

        debug!(
            "Quantizing tensor with shape {:?} using {:?}",
            tensor.shape, self.config.scheme
        );

        // Apply outlier handling if configured
        let processed_data = self.handle_outliers(&tensor.data).await?;

        // Compute quantization parameters
        let params = self.compute_quantization_params(&processed_data).await?;

        // Perform quantization based on scheme
        let quantized_data = match self.config.scheme {
            QuantizationScheme::INT8 => self.quantize_int8(&processed_data, &params).await?,
            QuantizationScheme::INT4 => self.quantize_int4(&processed_data, &params).await?,
            QuantizationScheme::FP8E4M3 => self.quantize_fp8_e4m3(&processed_data, &params).await?,
            QuantizationScheme::FP8E5M2 => self.quantize_fp8_e5m2(&processed_data, &params).await?,
        };

        let target_dtype = match self.config.scheme {
            QuantizationScheme::INT8 => DataType::INT8,
            QuantizationScheme::INT4 => DataType::INT4,
            QuantizationScheme::FP8E4M3 => DataType::FP8E4M3,
            QuantizationScheme::FP8E5M2 => DataType::FP8E5M2,
        };

        Ok(TestTensor {
            data: quantized_data,
            shape: tensor.shape.clone(),
            dtype: target_dtype,
        })
    }

    /// Dequantize a tensor
    pub async fn dequantize(&self, tensor: &TestTensor) -> InferenceResult<TestTensor> {
        tensor.validate()?;

        debug!(
            "Dequantizing tensor with shape {:?} from {:?}",
            tensor.shape, tensor.dtype
        );

        // Use stored parameters or compute them
        let params = if let Some(ref params) = self.params {
            params.clone()
        } else {
            // For dequantization, we need to infer parameters from the quantized data
            self.infer_quantization_params(tensor).await?
        };

        // Perform dequantization based on original scheme
        let dequantized_data = match tensor.dtype {
            DataType::INT8 => self.dequantize_int8(&tensor.data, &params).await?,
            DataType::INT4 => self.dequantize_int4(&tensor.data, &params).await?,
            DataType::FP8E4M3 => self.dequantize_fp8_e4m3(&tensor.data, &params).await?,
            DataType::FP8E5M2 => self.dequantize_fp8_e5m2(&tensor.data, &params).await?,
            _ => {
                return Err(InferenceError::invalid_request(format!(
                    "Cannot dequantize tensor with dtype {:?}",
                    tensor.dtype
                )));
            }
        };

        Ok(TestTensor {
            data: dequantized_data,
            shape: tensor.shape.clone(),
            dtype: DataType::F32,
        })
    }

    // Implementation methods

    async fn handle_outliers(&self, data: &[f32]) -> InferenceResult<Vec<f32>> {
        match &self.config.outlier_handling {
            OutlierHandling::None => Ok(data.to_vec()),
            OutlierHandling::Clipping { percentile } => {
                let mut sorted_data = data.to_vec();
                sorted_data.sort_by(f32::total_cmp);

                let lower_idx = ((100.0 - percentile) / 200.0 * sorted_data.len() as f32) as usize;
                let upper_idx = (percentile / 100.0 * sorted_data.len() as f32) as usize;

                let min_val = sorted_data[lower_idx];
                let max_val = sorted_data[upper_idx.min(sorted_data.len() - 1)];

                Ok(data.iter().map(|&x| x.clamp(min_val, max_val)).collect())
            }
            OutlierHandling::Separate { threshold: _ } => {
                // For now, implement as clipping at 3 standard deviations
                let mean = data.iter().sum::<f32>() / data.len() as f32;
                let variance =
                    data.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / data.len() as f32;
                let std_dev = variance.sqrt();

                let min_val = mean - 3.0 * std_dev;
                let max_val = mean + 3.0 * std_dev;

                Ok(data.iter().map(|&x| x.clamp(min_val, max_val)).collect())
            }
        }
    }

    async fn compute_quantization_params(
        &self,
        data: &[f32],
    ) -> InferenceResult<QuantizationParams> {
        let mut params = QuantizationParams::new(self.config.scheme);

        // Find min and max values
        let min_val = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
        let max_val = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));

        params.min_value = min_val;
        params.max_value = max_val;

        // Compute scale and zero point based on scaling method and scheme
        match self.config.scaling_method {
            ScalingMethod::Symmetric => {
                let abs_max = max_val.abs().max(min_val.abs());
                params.scale = match self.config.scheme {
                    QuantizationScheme::INT8 => abs_max / 127.0,
                    QuantizationScheme::INT4 => abs_max / 7.0,
                    QuantizationScheme::FP8E4M3 | QuantizationScheme::FP8E5M2 => abs_max / 240.0, // Approximate
                };
                params.zero_point = 0;
            }
            ScalingMethod::Asymmetric => {
                let (qmin, qmax) = match self.config.scheme {
                    QuantizationScheme::INT8 => (-128i32, 127i32),
                    QuantizationScheme::INT4 => (-8i32, 7i32),
                    QuantizationScheme::FP8E4M3 | QuantizationScheme::FP8E5M2 => (-240i32, 239i32), // Approximate
                };

                params.scale = (max_val - min_val) / (qmax - qmin) as f32;
                params.zero_point = qmin - (min_val / params.scale).round() as i32;
            }
        }

        // Ensure scale is not zero
        if params.scale == 0.0 {
            params.scale = 1e-8;
        }

        Ok(params)
    }

    async fn infer_quantization_params(
        &self,
        tensor: &TestTensor,
    ) -> InferenceResult<QuantizationParams> {
        // For dequantization, we need to infer reasonable parameters
        // This is a simplified approach - in practice, these would be stored with the quantized tensor
        let mut params = QuantizationParams::new(match tensor.dtype {
            DataType::INT8 => QuantizationScheme::INT8,
            DataType::INT4 => QuantizationScheme::INT4,
            DataType::FP8E4M3 => QuantizationScheme::FP8E4M3,
            DataType::FP8E5M2 => QuantizationScheme::FP8E5M2,
            _ => {
                return Err(InferenceError::internal_error(
                    "quantization",
                    "Invalid dtype for inference",
                ));
            }
        });

        // Use default scale factors for inference
        params.scale = match tensor.dtype {
            DataType::INT8 => 1.0 / 127.0,
            DataType::INT4 => 1.0 / 7.0,
            DataType::FP8E4M3 | DataType::FP8E5M2 => 1.0 / 240.0,
            _ => 1.0,
        };
        params.zero_point = 0;

        Ok(params)
    }

    async fn quantize_int8(
        &self,
        data: &[f32],
        params: &QuantizationParams,
    ) -> InferenceResult<Vec<f32>> {
        let quantized: Vec<f32> = data
            .iter()
            .map(|&x| {
                let q_val = (x / params.scale).round() + params.zero_point as f32;
                q_val.clamp(-128.0, 127.0)
            })
            .collect();
        Ok(quantized)
    }

    async fn dequantize_int8(
        &self,
        data: &[f32],
        params: &QuantizationParams,
    ) -> InferenceResult<Vec<f32>> {
        let dequantized: Vec<f32> = data
            .iter()
            .map(|&q| (q - params.zero_point as f32) * params.scale)
            .collect();
        Ok(dequantized)
    }

    async fn quantize_int4(
        &self,
        data: &[f32],
        params: &QuantizationParams,
    ) -> InferenceResult<Vec<f32>> {
        let quantized: Vec<f32> = data
            .iter()
            .map(|&x| {
                let q_val = (x / params.scale).round() + params.zero_point as f32;
                q_val.clamp(-8.0, 7.0)
            })
            .collect();
        Ok(quantized)
    }

    async fn dequantize_int4(
        &self,
        data: &[f32],
        params: &QuantizationParams,
    ) -> InferenceResult<Vec<f32>> {
        let dequantized: Vec<f32> = data
            .iter()
            .map(|&q| (q - params.zero_point as f32) * params.scale)
            .collect();
        Ok(dequantized)
    }

    async fn quantize_fp8_e4m3(
        &self,
        data: &[f32],
        params: &QuantizationParams,
    ) -> InferenceResult<Vec<f32>> {
        // Simplified FP8 E4M3 quantization (4-bit exponent, 3-bit mantissa)
        let quantized: Vec<f32> = data
            .iter()
            .map(|&x| {
                let scaled = x / params.scale;
                // Simulate FP8 precision by rounding to representable values
                let rounded = (scaled * 8.0).round() / 8.0; // Coarse approximation
                rounded.clamp(-240.0, 239.0)
            })
            .collect();
        Ok(quantized)
    }

    async fn dequantize_fp8_e4m3(
        &self,
        data: &[f32],
        params: &QuantizationParams,
    ) -> InferenceResult<Vec<f32>> {
        let dequantized: Vec<f32> = data.iter().map(|&q| q * params.scale).collect();
        Ok(dequantized)
    }

    async fn quantize_fp8_e5m2(
        &self,
        data: &[f32],
        params: &QuantizationParams,
    ) -> InferenceResult<Vec<f32>> {
        // Simplified FP8 E5M2 quantization (5-bit exponent, 2-bit mantissa)
        let quantized: Vec<f32> = data
            .iter()
            .map(|&x| {
                let scaled = x / params.scale;
                // Simulate FP8 precision with different rounding for E5M2
                let rounded = (scaled * 4.0).round() / 4.0; // Coarse approximation
                rounded.clamp(-57344.0, 57344.0) // Approximate E5M2 range
            })
            .collect();
        Ok(quantized)
    }

    async fn dequantize_fp8_e5m2(
        &self,
        data: &[f32],
        params: &QuantizationParams,
    ) -> InferenceResult<Vec<f32>> {
        let dequantized: Vec<f32> = data.iter().map(|&q| q * params.scale).collect();
        Ok(dequantized)
    }
}

/// Dynamic quantizer with calibration support
pub struct DynamicQuantizer {
    config: QuantizationConfig,
    calibration_data: Arc<Mutex<Vec<TestTensor>>>,
    is_calibrated: Arc<Mutex<bool>>,
    params: Arc<Mutex<Option<QuantizationParams>>>,
}

impl DynamicQuantizer {
    /// Create a new dynamic quantizer
    #[must_use]
    pub fn new(config: QuantizationConfig) -> Self {
        Self {
            config,
            calibration_data: Arc::new(Mutex::new(Vec::new())),
            is_calibrated: Arc::new(Mutex::new(false)),
            params: Arc::new(Mutex::new(None)),
        }
    }

    /// Perform calibration with provided data
    pub async fn calibrate(&mut self, data: &[TestTensor]) -> InferenceResult<()> {
        info!("Starting calibration with {} samples", data.len());

        // Collect calibration statistics
        let mut all_values = Vec::new();
        for tensor in data {
            tensor.validate()?;
            all_values.extend_from_slice(&tensor.data);
        }

        if all_values.is_empty() {
            return Err(InferenceError::invalid_request(
                "No calibration data provided",
            ));
        }

        // Compute optimal quantization parameters using specified method
        let params = match self.config.calibration_method {
            CalibrationMethod::MinMax => self.calibrate_minmax(&all_values).await?,
            CalibrationMethod::Percentile => self.calibrate_percentile(&all_values).await?,
            CalibrationMethod::KLDivergence => self.calibrate_kl_divergence(&all_values).await?,
            CalibrationMethod::MSE => self.calibrate_mse(&all_values).await?,
        };

        *self.params.lock().await = Some(params);
        *self.is_calibrated.lock().await = true;

        info!("Calibration completed successfully");
        Ok(())
    }

    /// Quantize using calibrated parameters
    pub async fn quantize(&self, tensor: &TestTensor) -> InferenceResult<TestTensor> {
        let is_calibrated = *self.is_calibrated.lock().await;
        if !is_calibrated {
            return Err(InferenceError::invalid_request(
                "Quantizer must be calibrated before use",
            ));
        }

        // Create a temporary quantizer with calibrated parameters
        let quantizer = Quantizer {
            config: self.config.clone(),
            params: self.params.lock().await.clone(),
        };

        quantizer.quantize(tensor).await
    }

    /// Dequantize using calibrated parameters
    pub async fn dequantize(&self, tensor: &TestTensor) -> InferenceResult<TestTensor> {
        let is_calibrated = *self.is_calibrated.lock().await;
        if !is_calibrated {
            return Err(InferenceError::invalid_request(
                "Quantizer must be calibrated before use",
            ));
        }

        let quantizer = Quantizer {
            config: self.config.clone(),
            params: self.params.lock().await.clone(),
        };

        quantizer.dequantize(tensor).await
    }

    // Calibration method implementations

    async fn calibrate_minmax(&self, data: &[f32]) -> InferenceResult<QuantizationParams> {
        let min_val = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
        let max_val = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));

        let mut params = QuantizationParams::new(self.config.scheme);
        params.min_value = min_val;
        params.max_value = max_val;

        let abs_max = max_val.abs().max(min_val.abs());
        params.scale = match self.config.scheme {
            QuantizationScheme::INT8 => abs_max / 127.0,
            QuantizationScheme::INT4 => abs_max / 7.0,
            QuantizationScheme::FP8E4M3 | QuantizationScheme::FP8E5M2 => abs_max / 240.0,
        };
        params.zero_point = 0;

        if params.scale == 0.0 {
            params.scale = 1e-8;
        }

        Ok(params)
    }

    async fn calibrate_percentile(&self, data: &[f32]) -> InferenceResult<QuantizationParams> {
        let mut sorted_data = data.to_vec();
        sorted_data.sort_by(f32::total_cmp);

        // Use 99.9th percentile to handle outliers
        let percentile = 99.9;
        let lower_idx = ((100.0 - percentile) / 200.0 * sorted_data.len() as f32) as usize;
        let upper_idx = (percentile / 100.0 * sorted_data.len() as f32) as usize;

        let min_val = sorted_data[lower_idx];
        let max_val = sorted_data[upper_idx.min(sorted_data.len() - 1)];

        let mut params = QuantizationParams::new(self.config.scheme);
        params.min_value = min_val;
        params.max_value = max_val;

        let abs_max = max_val.abs().max(min_val.abs());
        params.scale = match self.config.scheme {
            QuantizationScheme::INT8 => abs_max / 127.0,
            QuantizationScheme::INT4 => abs_max / 7.0,
            QuantizationScheme::FP8E4M3 | QuantizationScheme::FP8E5M2 => abs_max / 240.0,
        };
        params.zero_point = 0;

        if params.scale == 0.0 {
            params.scale = 1e-8;
        }

        Ok(params)
    }

    async fn calibrate_kl_divergence(&self, data: &[f32]) -> InferenceResult<QuantizationParams> {
        // Simplified KL divergence calibration
        // In practice, this would involve computing histogram distributions
        // and finding the threshold that minimizes KL divergence

        // For now, use a percentile-based approach as an approximation
        self.calibrate_percentile(data).await
    }

    async fn calibrate_mse(&self, data: &[f32]) -> InferenceResult<QuantizationParams> {
        // Simplified MSE-based calibration
        // This would involve trying different scales and finding the one with minimum MSE

        let mut best_params = self.calibrate_minmax(data).await?;
        let mut best_mse = f32::INFINITY;

        // Try different scale factors
        let base_params = best_params.clone();
        for factor in [0.8, 0.9, 1.0, 1.1, 1.2] {
            let mut test_params = base_params.clone();
            test_params.scale *= factor;

            // Compute MSE for this scale
            let mut mse = 0.0;
            for &value in data {
                let quantized = (value / test_params.scale).round() * test_params.scale;
                mse += (value - quantized).powi(2);
            }
            mse /= data.len() as f32;

            if mse < best_mse {
                best_mse = mse;
                best_params = test_params;
            }
        }

        Ok(best_params)
    }
}

/// Validation metrics for quantization quality
#[derive(Debug, Clone)]
pub struct ValidationMetrics {
    pub mse: f32,
    pub snr: f32,
    pub relative_error: f32,
    pub max_error: f32,
    pub threshold: f32,
}

impl ValidationMetrics {
    #[must_use]
    pub fn passes_threshold(&self) -> bool {
        self.mse < self.threshold
    }
}

/// Validator for quantization accuracy
pub struct QuantizationValidator {
    config: QuantizationConfig,
}

impl QuantizationValidator {
    #[must_use]
    pub fn new(config: QuantizationConfig) -> Self {
        Self { config }
    }

    pub async fn validate(
        &self,
        original: &TestTensor,
        quantized: &TestTensor,
    ) -> InferenceResult<ValidationMetrics> {
        if original.data.len() != quantized.data.len() {
            return Err(InferenceError::invalid_request("Tensor size mismatch"));
        }

        let mse = self.calculate_mse(&original.data, &quantized.data);
        let snr = self.calculate_snr(&original.data, &quantized.data);
        let relative_error = self.calculate_relative_error(&original.data, &quantized.data);
        let max_error = self.calculate_max_error(&original.data, &quantized.data);

        Ok(ValidationMetrics {
            mse,
            snr,
            relative_error,
            max_error,
            threshold: self.config.accuracy_threshold,
        })
    }

    fn calculate_mse(&self, original: &[f32], quantized: &[f32]) -> f32 {
        let sum_squared_error: f32 = original
            .iter()
            .zip(quantized.iter())
            .map(|(o, q)| (o - q) * (o - q))
            .sum();
        sum_squared_error / original.len() as f32
    }

    fn calculate_snr(&self, original: &[f32], quantized: &[f32]) -> f32 {
        let signal_power: f32 = original.iter().map(|x| x * x).sum();
        let noise_power: f32 = original
            .iter()
            .zip(quantized.iter())
            .map(|(o, q)| (o - q) * (o - q))
            .sum();

        if noise_power == 0.0 {
            f32::INFINITY
        } else {
            10.0 * (signal_power / noise_power).log10()
        }
    }

    fn calculate_relative_error(&self, original: &[f32], quantized: &[f32]) -> f32 {
        let sum_relative_error: f32 = original
            .iter()
            .zip(quantized.iter())
            .map(|(o, q)| {
                if o.abs() < 1e-8 {
                    if q.abs() < 1e-8 { 0.0 } else { 1.0 }
                } else {
                    ((o - q) / o).abs()
                }
            })
            .sum();
        sum_relative_error / original.len() as f32
    }

    fn calculate_max_error(&self, original: &[f32], quantized: &[f32]) -> f32 {
        original
            .iter()
            .zip(quantized.iter())
            .map(|(o, q)| (o - q).abs())
            .fold(0.0f32, f32::max)
    }
}

/// Configuration for individual layers in mixed precision
#[derive(Debug, Clone)]
pub struct LayerConfig {
    pub scheme: QuantizationScheme,
    pub config: QuantizationConfig,
}

impl LayerConfig {
    #[must_use]
    pub fn new(scheme: QuantizationScheme) -> Self {
        Self {
            scheme,
            config: QuantizationConfig::new().with_scheme(scheme),
        }
    }
}

/// Mixed precision quantizer for different model layers
pub struct MixedPrecisionQuantizer {
    layer_configs: HashMap<String, LayerConfig>,
    quantizers: HashMap<String, Quantizer>,
}

impl Default for MixedPrecisionQuantizer {
    fn default() -> Self {
        Self::new()
    }
}

impl MixedPrecisionQuantizer {
    #[must_use]
    pub fn new() -> Self {
        Self {
            layer_configs: HashMap::new(),
            quantizers: HashMap::new(),
        }
    }

    pub fn add_layer_config(&mut self, layer_name: &str, config: LayerConfig) {
        let quantizer = Quantizer::new(config.config.clone());
        self.layer_configs.insert(layer_name.to_string(), config);
        self.quantizers.insert(layer_name.to_string(), quantizer);
    }

    pub async fn quantize_layer(
        &self,
        layer_name: &str,
        tensor: &TestTensor,
    ) -> InferenceResult<TestTensor> {
        let quantizer = self.quantizers.get(layer_name).ok_or_else(|| {
            InferenceError::invalid_request(format!("No configuration for layer: {layer_name}"))
        })?;

        quantizer.quantize(tensor).await
    }

    pub async fn dequantize_layer(
        &self,
        layer_name: &str,
        tensor: &TestTensor,
    ) -> InferenceResult<TestTensor> {
        let quantizer = self.quantizers.get(layer_name).ok_or_else(|| {
            InferenceError::invalid_request(format!("No configuration for layer: {layer_name}"))
        })?;

        quantizer.dequantize(tensor).await
    }
}

/// Batch quantizer for processing multiple tensors
pub struct BatchQuantizer {
    quantizer: Quantizer,
}

impl BatchQuantizer {
    #[must_use]
    pub fn new(config: QuantizationConfig) -> Self {
        Self {
            quantizer: Quantizer::new(config),
        }
    }

    pub async fn quantize_batch(&self, tensors: &[TestTensor]) -> InferenceResult<Vec<TestTensor>> {
        let mut results = Vec::with_capacity(tensors.len());
        for tensor in tensors {
            results.push(self.quantizer.quantize(tensor).await?);
        }
        Ok(results)
    }

    pub async fn dequantize_batch(
        &self,
        tensors: &[TestTensor],
    ) -> InferenceResult<Vec<TestTensor>> {
        let mut results = Vec::with_capacity(tensors.len());
        for tensor in tensors {
            results.push(self.quantizer.dequantize(tensor).await?);
        }
        Ok(results)
    }
}

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

    #[test]
    fn test_config_validation() {
        let valid_config = QuantizationConfig::new();
        assert!(valid_config.validate().is_ok());

        let invalid_config = QuantizationConfig::new().with_calibration_samples(0);
        assert!(invalid_config.validate().is_err());
    }

    #[test]
    fn test_tensor_creation() {
        let tensor = TestTensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2]);
        assert_eq!(tensor.data.len(), 4);
        assert_eq!(tensor.shape, vec![2, 2]);
        assert_eq!(tensor.num_elements(), 4);
        assert_eq!(tensor.dtype, DataType::F32);
    }

    #[tokio::test]
    async fn test_basic_quantization() {
        let config = QuantizationConfig::new();
        let quantizer = Quantizer::new(config);

        let tensor = TestTensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![4]);
        let quantized = quantizer.quantize(&tensor).await.unwrap();
        let dequantized = quantizer.dequantize(&quantized).await.unwrap();

        assert_eq!(quantized.dtype, DataType::INT8);
        assert_eq!(dequantized.dtype, DataType::F32);
        assert_eq!(dequantized.data.len(), tensor.data.len());
    }
}