mlmf 0.2.0

//! Advanced quantization support for ML models.
//!
//! This module provides comprehensive quantization functionality including:
//!
//! - Post-training quantization (PTQ) with calibration data
//! - Various quantization schemes (INT8, INT4, mixed precision)
//! - Dynamic and static quantization modes
//! - Quantization-aware loading and inference
//! - Calibration dataset handling
//! - Quantization profiling and optimization

use crate::progress::{ProgressEvent, ProgressFn};
use crate::{Error, LoadOptions, LoadedModel};
use candle_core::{DType, Device, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;

/// Quantization data types
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum QuantizationType {
    /// 8-bit integer quantization
    Int8,
    /// 4-bit integer quantization  
    Int4,
    /// Mixed precision (different layers use different precisions)
    Mixed,
    /// Dynamic quantization (quantize weights only)
    Dynamic,
    /// Static quantization (quantize weights and activations)
    Static,
}

impl Default for QuantizationType {
    fn default() -> Self {
        Self::Int8
    }
}

/// Quantization scheme for a layer or tensor
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuantizationScheme {
    /// Quantization type for this layer
    pub quant_type: QuantizationType,
    /// Scale factor for quantization
    pub scale: f32,
    /// Zero point for asymmetric quantization
    pub zero_point: i32,
    /// Whether to use symmetric or asymmetric quantization
    pub symmetric: bool,
    /// Quantization range (min, max)
    pub range: (f32, f32),
}

impl Default for QuantizationScheme {
    fn default() -> Self {
        Self {
            quant_type: QuantizationType::Int8,
            scale: 1.0,
            zero_point: 0,
            symmetric: true,
            range: (-128.0, 127.0),
        }
    }
}

/// Configuration for quantization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuantizationConfig {
    /// Default quantization type
    pub quantization_type: QuantizationType,
    /// Number of calibration samples to use
    pub calibration_samples: usize,
    /// Calibration method ("minmax", "entropy", "percentile", "kl_divergence")
    pub calibration_method: String,
    /// Percentile for percentile-based calibration (0.0-100.0)
    pub percentile: f32,
    /// Whether to use symmetric quantization
    pub symmetric: bool,
    /// Layer-specific quantization overrides
    pub layer_config: HashMap<String, QuantizationScheme>,
    /// Layers to skip quantization
    pub skip_layers: Vec<String>,
    /// Whether to quantize bias terms
    pub quantize_bias: bool,
    /// Whether to use block-wise quantization
    pub block_wise: bool,
    /// Block size for block-wise quantization
    pub block_size: usize,
    /// Enable advanced activation statistics collection
    pub advanced_stats: bool,
    /// Entropy bins for entropy-based calibration
    pub entropy_bins: usize,
    /// KL divergence threshold for optimal quantization
    pub kl_threshold: f32,
}

impl Default for QuantizationConfig {
    fn default() -> Self {
        Self {
            quantization_type: QuantizationType::Int8,
            calibration_samples: 128,
            calibration_method: "minmax".to_string(),
            percentile: 99.9,
            symmetric: true,
            layer_config: HashMap::new(),
            skip_layers: vec!["embedding".to_string(), "norm".to_string()],
            quantize_bias: false,
            block_wise: false,
            block_size: 128,
            advanced_stats: false,
            entropy_bins: 2048,
            kl_threshold: 0.1,
        }
    }
}

/// Calibration dataset for quantization
#[derive(Debug)]
pub struct CalibrationDataset {
    /// Input tensors for calibration
    pub samples: Vec<HashMap<String, Tensor>>,
    /// Device where tensors are stored
    pub device: Device,
}

impl CalibrationDataset {
    /// Create new calibration dataset
    pub fn new(device: Device) -> Self {
        Self {
            samples: Vec::new(),
            device,
        }
    }

    /// Add a calibration sample
    pub fn add_sample(&mut self, sample: HashMap<String, Tensor>) -> Result<(), Error> {
        // For simplicity, assume tensors are already on correct device
        self.samples.push(sample);
        Ok(())
    }

    /// Get number of samples
    pub fn len(&self) -> usize {
        self.samples.len()
    }

    /// Check if dataset is empty
    pub fn is_empty(&self) -> bool {
        self.samples.is_empty()
    }

    /// Get a sample by index
    pub fn get_sample(&self, index: usize) -> Option<&HashMap<String, Tensor>> {
        self.samples.get(index)
    }
}

/// Statistics collected during calibration
#[derive(Debug, Clone)]
pub struct ActivationStats {
    /// Minimum values observed
    pub min_val: f32,
    /// Maximum values observed
    pub max_val: f32,
    /// Mean values
    pub mean_val: f32,
    /// Standard deviation
    pub std_val: f32,
    /// Number of samples processed
    pub num_samples: usize,
    /// Histogram bins for entropy calculation
    pub histogram: Option<Vec<f32>>,
    /// Percentile values (P1, P5, P95, P99)
    pub percentiles: Option<Vec<f32>>,
    /// KL divergence scores for different quantization schemes
    pub kl_scores: Option<HashMap<String, f32>>,
}

/// Extended quantization context for advanced features
#[derive(Debug)]
pub struct QuantizationContext {
    /// External metadata storage
    pub metadata: HashMap<String, serde_json::Value>,
    /// Extended tensor statistics
    pub tensor_stats: HashMap<String, ActivationStats>,
    /// Layer-specific quantization schemes (overrides config)
    pub layer_schemes: HashMap<String, QuantizationScheme>,
    /// Block-wise quantization blocks per tensor
    pub tensor_blocks: HashMap<String, Vec<QuantizationScheme>>,
    /// Performance metrics during quantization
    pub metrics: QuantizationMetrics,
}

/// Performance metrics for quantization
#[derive(Debug, Default)]
pub struct QuantizationMetrics {
    /// Compression ratio achieved
    pub compression_ratio: f32,
    /// Average quantization error
    pub avg_error: f32,
    /// Peak signal-to-noise ratio
    pub psnr: f32,
    /// Per-layer compression ratios
    pub layer_compression: HashMap<String, f32>,
    /// Calibration time in seconds
    pub calibration_time: f64,
    /// Quantization time in seconds
    pub quantization_time: f64,
}

impl QuantizationContext {
    /// Create new empty context
    pub fn new() -> Self {
        Self {
            metadata: HashMap::new(),
            tensor_stats: HashMap::new(),
            layer_schemes: HashMap::new(),
            tensor_blocks: HashMap::new(),
            metrics: QuantizationMetrics::default(),
        }
    }

    /// Add metadata entry
    pub fn add_metadata(&mut self, key: String, value: serde_json::Value) {
        self.metadata.insert(key, value);
    }

    /// Get metadata entry
    pub fn get_metadata(&self, key: &str) -> Option<&serde_json::Value> {
        self.metadata.get(key)
    }

    /// Set layer-specific quantization scheme
    pub fn set_layer_scheme(&mut self, layer_name: String, scheme: QuantizationScheme) {
        self.layer_schemes.insert(layer_name, scheme);
    }

    /// Get layer-specific quantization scheme
    pub fn get_layer_scheme(&self, layer_name: &str) -> Option<&QuantizationScheme> {
        self.layer_schemes.get(layer_name)
    }
}

/// Quantization engine for performing model quantization
#[derive(Debug)]
pub struct QuantizationEngine {
    /// Quantization configuration
    config: QuantizationConfig,
    /// Target device
    device: Device,
    /// Activation statistics per layer
    activation_stats: HashMap<String, ActivationStats>,
}

impl QuantizationEngine {
    /// Create new quantization engine
    pub fn new(config: QuantizationConfig, device: Device) -> Self {
        Self {
            config,
            device,
            activation_stats: HashMap::new(),
        }
    }

    /// Perform calibration on a model using calibration dataset
    pub fn calibrate(
        &mut self,
        model: &LoadedModel,
        calibration_data: &CalibrationDataset,
        progress_callback: Option<ProgressFn>,
    ) -> Result<(), Error> {
        if let Some(ref callback) = progress_callback {
            callback(ProgressEvent::Status {
                message: "Starting model calibration for quantization".to_string(),
            });
        }

        let num_samples = std::cmp::min(calibration_data.len(), self.config.calibration_samples);

        // Collect statistics from model tensors (simplified calibration)
        for tensor_name in model.raw_tensors.keys() {
            if self.should_skip_layer(tensor_name) {
                continue;
            }

            if let Some(tensor) = model.raw_tensors.get(tensor_name) {
                let stats = self.collect_tensor_stats(tensor)?;
                self.activation_stats.insert(tensor_name.clone(), stats);
            }
        }

        if let Some(ref callback) = progress_callback {
            callback(ProgressEvent::Status {
                message: "Calibration completed".to_string(),
            });
        }

        Ok(())
    }

    /// Collect statistics from a tensor
    fn collect_tensor_stats(&self, tensor: &Tensor) -> Result<ActivationStats, Error> {
        // Basic statistics computation
        let min_val = tensor
            .min_keepdim(0)?
            .to_scalar::<f32>()
            .map_err(|e| Error::Candle(e))?;
        let max_val = tensor
            .max_keepdim(0)?
            .to_scalar::<f32>()
            .map_err(|e| Error::Candle(e))?;
        let mean_val = tensor
            .mean_all()?
            .to_scalar::<f32>()
            .map_err(|e| Error::Candle(e))?;

        // Standard deviation calculation
        let mean_tensor = Tensor::new(&[mean_val], &Device::Cpu)?.to_device(&tensor.device())?;
        let diff = tensor.broadcast_sub(&mean_tensor)?;
        let variance = diff
            .powf(2.0)?
            .mean_all()?
            .to_scalar::<f32>()
            .map_err(|e| Error::Candle(e))?;
        let std_val = variance.sqrt();

        // Advanced statistics if enabled
        let (histogram, percentiles, kl_scores) = if self.config.advanced_stats {
            let hist = self.compute_histogram(tensor)?;
            let percs = self.compute_percentiles(tensor)?;
            let kl = self.compute_kl_scores(tensor)?;
            (Some(hist), Some(percs), Some(kl))
        } else {
            (None, None, None)
        };

        Ok(ActivationStats {
            min_val,
            max_val,
            mean_val,
            std_val,
            num_samples: 1,
            histogram,
            percentiles,
            kl_scores,
        })
    }

    /// Compute histogram for entropy calculation
    fn compute_histogram(&self, tensor: &Tensor) -> Result<Vec<f32>, Error> {
        let flat = tensor.flatten_all().map_err(|e| Error::Candle(e))?;
        let data: Vec<f32> = flat.to_vec1().map_err(|e| Error::Candle(e))?;

        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 histogram = vec![0.0; self.config.entropy_bins];
        let bin_size = (max_val - min_val) / (self.config.entropy_bins as f32);

        if bin_size > 0.0 {
            for &value in &data {
                let bin = ((value - min_val) / bin_size).floor() as usize;
                let bin = bin.min(self.config.entropy_bins - 1);
                histogram[bin] += 1.0;
            }

            // Normalize histogram
            let total = data.len() as f32;
            for bin in &mut histogram {
                *bin /= total;
            }
        }

        Ok(histogram)
    }

    /// Compute percentile values
    fn compute_percentiles(&self, tensor: &Tensor) -> Result<Vec<f32>, Error> {
        let flat = tensor.flatten_all().map_err(|e| Error::Candle(e))?;
        let mut data: Vec<f32> = flat.to_vec1().map_err(|e| Error::Candle(e))?;
        data.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let percentiles = vec![1.0, 5.0, 95.0, 99.0];
        let mut results = Vec::new();

        for &p in &percentiles {
            let index = ((p / 100.0) * (data.len() - 1) as f32) as usize;
            results.push(data.get(index).copied().unwrap_or(0.0));
        }

        Ok(results)
    }

    /// Compute KL divergence scores for different quantization schemes
    fn compute_kl_scores(&self, tensor: &Tensor) -> Result<HashMap<String, f32>, Error> {
        let mut scores = HashMap::new();

        // Compute KL divergence for INT8 and INT4 quantization
        for &bits in &[4, 8] {
            let scheme_name = format!("int{}", bits);
            let score = self.compute_kl_divergence_for_bits(tensor, bits)?;
            scores.insert(scheme_name, score);
        }

        Ok(scores)
    }

    /// Compute KL divergence for specific bit width
    fn compute_kl_divergence_for_bits(&self, tensor: &Tensor, bits: u8) -> Result<f32, Error> {
        // Simplified KL divergence calculation
        // In practice, this would compare original vs quantized distributions
        let flat = tensor.flatten_all().map_err(|e| Error::Candle(e))?;
        let data: Vec<f32> = flat.to_vec1().map_err(|e| Error::Candle(e))?;

        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 levels = (1 << bits) as f32;
        let scale = (max_val - min_val) / levels;

        // Simple approximation: smaller scale = better quantization = lower KL divergence
        Ok(scale.abs())
    }

    /// Check if a layer should be skipped during quantization
    fn should_skip_layer(&self, layer_name: &str) -> bool {
        self.config
            .skip_layers
            .iter()
            .any(|skip| layer_name.contains(skip))
    }

    /// Quantize model tensors
    pub fn quantize_model(
        &self,
        model: &LoadedModel,
        progress_callback: Option<ProgressFn>,
    ) -> Result<HashMap<String, Tensor>, Error> {
        if let Some(ref callback) = progress_callback {
            callback(ProgressEvent::Status {
                message: "Starting model quantization".to_string(),
            });
        }

        let mut quantized_tensors = HashMap::new();
        let total_tensors = model.raw_tensors.len();

        for (idx, (tensor_name, tensor)) in model.raw_tensors.iter().enumerate() {
            if let Some(ref callback) = progress_callback {
                if idx % 10 == 0 {
                    callback(ProgressEvent::Status {
                        message: format!(
                            "Quantizing tensor {}/{}: {}",
                            idx + 1,
                            total_tensors,
                            tensor_name
                        ),
                    });
                }
            }

            let quantized_tensor = if self.should_skip_layer(tensor_name) {
                // Keep original tensor for skipped layers
                tensor.clone()
            } else {
                // Get quantization scheme for this layer
                let scheme = self.get_quantization_scheme(tensor_name, tensor)?;
                self.quantize_tensor(tensor, &scheme)?
            };

            quantized_tensors.insert(tensor_name.clone(), quantized_tensor);
        }

        if let Some(ref callback) = progress_callback {
            callback(ProgressEvent::Status {
                message: "Model quantization completed".to_string(),
            });
        }

        Ok(quantized_tensors)
    }

    /// Get quantization scheme for a specific tensor
    fn get_quantization_scheme(
        &self,
        tensor_name: &str,
        tensor: &Tensor,
    ) -> Result<QuantizationScheme, Error> {
        // Check for layer-specific configuration
        if let Some(scheme) = self.config.layer_config.get(tensor_name) {
            return Ok(scheme.clone());
        }

        // Use calibration statistics if available
        if let Some(stats) = self.activation_stats.get(tensor_name) {
            return self.compute_scheme_from_stats(stats);
        }

        // Fallback: compute scheme from tensor directly
        self.compute_scheme_from_tensor(tensor)
    }

    /// Compute quantization scheme from calibration statistics
    fn compute_scheme_from_stats(
        &self,
        stats: &ActivationStats,
    ) -> Result<QuantizationScheme, Error> {
        let mut scheme = QuantizationScheme::default();
        scheme.quant_type = self.config.quantization_type;

        match self.config.calibration_method.as_str() {
            "minmax" => {
                scheme.range = (stats.min_val, stats.max_val);
            }
            "percentile" => {
                if let Some(ref percentiles) = stats.percentiles {
                    // Use actual computed percentiles
                    let p1 = percentiles.get(0).copied().unwrap_or(stats.min_val);
                    let p99 = percentiles.get(3).copied().unwrap_or(stats.max_val);
                    scheme.range = (p1, p99);
                } else {
                    // Fallback to standard deviation based
                    let factor = self.config.percentile / 100.0 * 3.0;
                    scheme.range = (
                        stats.mean_val - factor * stats.std_val,
                        stats.mean_val + factor * stats.std_val,
                    );
                }
            }
            "entropy" => {
                // Entropy-based calibration using histogram
                if let Some(ref histogram) = stats.histogram {
                    scheme.range = self.compute_entropy_range(histogram, stats)?;
                } else {
                    // Fallback to minmax
                    scheme.range = (stats.min_val, stats.max_val);
                }
            }
            "kl_divergence" => {
                // KL divergence-based calibration
                if let Some(ref kl_scores) = stats.kl_scores {
                    scheme.range = self.compute_kl_optimal_range(kl_scores, stats)?;
                } else {
                    // Fallback to minmax
                    scheme.range = (stats.min_val, stats.max_val);
                }
            }
            _ => {
                return Err(Error::invalid_config(format!(
                    "Unknown calibration method: {}. Supported: minmax, percentile, entropy, kl_divergence",
                    self.config.calibration_method
                )));
            }
        }

        self.compute_scale_and_zero_point(&mut scheme)?;
        Ok(scheme)
    }

    /// Compute optimal range using entropy-based calibration
    fn compute_entropy_range(
        &self,
        histogram: &[f32],
        stats: &ActivationStats,
    ) -> Result<(f32, f32), Error> {
        // Find range that minimizes entropy while preserving information
        let mut best_range = (stats.min_val, stats.max_val);
        let mut min_entropy = f32::INFINITY;

        let total_range = stats.max_val - stats.min_val;

        // Try different percentages of the range
        for percentage in [0.95, 0.98, 0.99, 0.995, 1.0] {
            let range_size = total_range * percentage;
            let center = stats.mean_val;
            let min_val = center - range_size / 2.0;
            let max_val = center + range_size / 2.0;

            // Compute entropy for this range
            let entropy = self.compute_range_entropy(histogram, min_val, max_val, stats)?;

            if entropy < min_entropy {
                min_entropy = entropy;
                best_range = (min_val.max(stats.min_val), max_val.min(stats.max_val));
            }
        }

        Ok(best_range)
    }

    /// Compute entropy for a specific range
    fn compute_range_entropy(
        &self,
        _histogram: &[f32],
        _min_val: f32,
        _max_val: f32,
        _stats: &ActivationStats,
    ) -> Result<f32, Error> {
        // Simplified entropy computation
        // In practice, this would compute the entropy of the quantized distribution
        // For now, return a placeholder value
        Ok(1.0)
    }

    /// Compute optimal range using KL divergence
    fn compute_kl_optimal_range(
        &self,
        kl_scores: &HashMap<String, f32>,
        stats: &ActivationStats,
    ) -> Result<(f32, f32), Error> {
        // Find the quantization scheme with minimum KL divergence
        let scheme_type = format!(
            "int{}",
            match self.config.quantization_type {
                QuantizationType::Int8 => 8,
                QuantizationType::Int4 => 4,
                _ => 8,
            }
        );

        if let Some(&kl_score) = kl_scores.get(&scheme_type) {
            // Use KL score to adjust range - lower score means better quantization
            let adjustment_factor = (kl_score * self.config.kl_threshold).min(0.1);
            let range = stats.max_val - stats.min_val;
            let adjusted_range = range * (1.0 - adjustment_factor);
            let center = (stats.max_val + stats.min_val) / 2.0;

            Ok((center - adjusted_range / 2.0, center + adjusted_range / 2.0))
        } else {
            // Fallback to minmax
            Ok((stats.min_val, stats.max_val))
        }
    }

    /// Compute quantization scheme directly from tensor
    fn compute_scheme_from_tensor(&self, tensor: &Tensor) -> Result<QuantizationScheme, Error> {
        let mut scheme = QuantizationScheme::default();
        scheme.quant_type = self.config.quantization_type;

        let min_val = tensor
            .min_keepdim(0)?
            .to_scalar::<f32>()
            .map_err(|e| Error::Candle(e))?;
        let max_val = tensor
            .max_keepdim(0)?
            .to_scalar::<f32>()
            .map_err(|e| Error::Candle(e))?;
        scheme.range = (min_val, max_val);

        self.compute_scale_and_zero_point(&mut scheme)?;
        Ok(scheme)
    }

    /// Compute scale and zero point for quantization scheme
    fn compute_scale_and_zero_point(&self, scheme: &mut QuantizationScheme) -> Result<(), Error> {
        let (min_val, max_val) = scheme.range;

        let (qmin, qmax) = match scheme.quant_type {
            QuantizationType::Int8 => (-128i32, 127i32),
            QuantizationType::Int4 => (-8i32, 7i32),
            _ => (-128i32, 127i32),
        };

        if self.config.symmetric {
            // Symmetric quantization
            let abs_max = min_val.abs().max(max_val.abs());
            scheme.scale = abs_max / qmax as f32;
            scheme.zero_point = 0;
            scheme.range = (-abs_max, abs_max);
        } else {
            // Asymmetric quantization
            scheme.scale = (max_val - min_val) / (qmax - qmin) as f32;
            scheme.zero_point = qmin - (min_val / scheme.scale).round() as i32;
        }

        Ok(())
    }

    /// Quantize a single tensor
    fn quantize_tensor(
        &self,
        tensor: &Tensor,
        scheme: &QuantizationScheme,
    ) -> Result<Tensor, Error> {
        if self.config.block_wise {
            self.quantize_tensor_blockwise(tensor, scheme)
        } else {
            match scheme.quant_type {
                QuantizationType::Int8 => self.quantize_tensor_int8(tensor, scheme),
                QuantizationType::Int4 => self.quantize_tensor_int4(tensor, scheme),
                QuantizationType::Dynamic => self.quantize_tensor_dynamic(tensor, scheme),
                _ => Ok(tensor.clone()), // Fallback for unsupported types
            }
        }
    }

    /// Quantize tensor using block-wise quantization
    fn quantize_tensor_blockwise(
        &self,
        tensor: &Tensor,
        scheme: &QuantizationScheme,
    ) -> Result<Tensor, Error> {
        let shape = tensor.shape().clone();
        let total_elements = shape.elem_count();

        if total_elements <= self.config.block_size {
            // Tensor is smaller than block size, use regular quantization
            return match scheme.quant_type {
                QuantizationType::Int8 => self.quantize_tensor_int8(tensor, scheme),
                QuantizationType::Int4 => self.quantize_tensor_int4(tensor, scheme),
                QuantizationType::Dynamic => self.quantize_tensor_dynamic(tensor, scheme),
                _ => Ok(tensor.clone()),
            };
        }

        // Flatten tensor for block processing
        let flat_tensor = tensor.flatten_all().map_err(|e| Error::Candle(e))?;
        let data: Vec<f32> = flat_tensor.to_vec1().map_err(|e| Error::Candle(e))?;

        let mut quantized_data = Vec::new();
        let num_blocks = (data.len() + self.config.block_size - 1) / self.config.block_size;

        for block_idx in 0..num_blocks {
            let start_idx = block_idx * self.config.block_size;
            let end_idx = (start_idx + self.config.block_size).min(data.len());
            let block_data = &data[start_idx..end_idx];

            // Create tensor for this block
            let block_tensor =
                Tensor::new(block_data, tensor.device()).map_err(|e| Error::Candle(e))?;

            // Compute block-specific quantization scheme
            let block_scheme = self.compute_block_quantization_scheme(&block_tensor, scheme)?;

            // Quantize the block
            let quantized_block = match scheme.quant_type {
                QuantizationType::Int8 => {
                    self.quantize_tensor_int8(&block_tensor, &block_scheme)?
                }
                QuantizationType::Int4 => {
                    self.quantize_tensor_int4(&block_tensor, &block_scheme)?
                }
                QuantizationType::Dynamic => {
                    self.quantize_tensor_dynamic(&block_tensor, &block_scheme)?
                }
                _ => block_tensor,
            };

            // Add quantized block data to result
            let block_result: Vec<f32> = quantized_block.to_vec1().map_err(|e| Error::Candle(e))?;
            quantized_data.extend_from_slice(&block_result);
        }

        // Reshape back to original shape
        let result_tensor = Tensor::new(&quantized_data[..data.len()], tensor.device())
            .map_err(|e| Error::Candle(e))?
            .reshape(&shape)
            .map_err(|e| Error::Candle(e))?;

        Ok(result_tensor)
    }

    /// Compute quantization scheme for a specific block
    fn compute_block_quantization_scheme(
        &self,
        block_tensor: &Tensor,
        base_scheme: &QuantizationScheme,
    ) -> Result<QuantizationScheme, Error> {
        let mut block_scheme = base_scheme.clone();

        // Compute block-specific statistics
        let min_val = block_tensor
            .min_keepdim(0)?
            .to_scalar::<f32>()
            .map_err(|e| Error::Candle(e))?;
        let max_val = block_tensor
            .max_keepdim(0)?
            .to_scalar::<f32>()
            .map_err(|e| Error::Candle(e))?;

        // Update range for this block
        block_scheme.range = (min_val, max_val);

        // Recompute scale and zero point for the block
        self.compute_scale_and_zero_point(&mut block_scheme)?;

        Ok(block_scheme)
    }

    /// Quantize tensor to INT8
    fn quantize_tensor_int8(
        &self,
        tensor: &Tensor,
        scheme: &QuantizationScheme,
    ) -> Result<Tensor, Error> {
        // Create scalar tensors for operations
        let scale_tensor =
            Tensor::new(&[scheme.scale], &Device::Cpu)?.to_device(&tensor.device())?;
        let zero_point_tensor =
            Tensor::new(&[scheme.zero_point as f32], &Device::Cpu)?.to_device(&tensor.device())?;

        // Quantize: q = round(x / scale + zero_point)
        let scaled = tensor
            .broadcast_div(&scale_tensor)
            .map_err(|e| Error::Candle(e))?;
        let shifted = scaled
            .broadcast_add(&zero_point_tensor)
            .map_err(|e| Error::Candle(e))?;
        let quantized = shifted.round().map_err(|e| Error::Candle(e))?;

        // Clamp to quantization range
        let clamped = quantized
            .clamp(-128.0, 127.0)
            .map_err(|e| Error::Candle(e))?;

        // Dequantize for storage: x = (q - zero_point) * scale
        let dequantized_shifted = clamped
            .broadcast_sub(&zero_point_tensor)
            .map_err(|e| Error::Candle(e))?;
        let dequantized = dequantized_shifted
            .broadcast_mul(&scale_tensor)
            .map_err(|e| Error::Candle(e))?;

        Ok(dequantized)
    }

    /// Quantize tensor to INT4 (simplified implementation)
    fn quantize_tensor_int4(
        &self,
        tensor: &Tensor,
        scheme: &QuantizationScheme,
    ) -> Result<Tensor, Error> {
        // Create scalar tensors for operations
        let scale_tensor =
            Tensor::new(&[scheme.scale], &Device::Cpu)?.to_device(&tensor.device())?;
        let zero_point_tensor =
            Tensor::new(&[scheme.zero_point as f32], &Device::Cpu)?.to_device(&tensor.device())?;

        let scaled = tensor
            .broadcast_div(&scale_tensor)
            .map_err(|e| Error::Candle(e))?;
        let shifted = scaled
            .broadcast_add(&zero_point_tensor)
            .map_err(|e| Error::Candle(e))?;
        let quantized = shifted.round().map_err(|e| Error::Candle(e))?;

        // Clamp to INT4 range
        let clamped = quantized.clamp(-8.0, 7.0).map_err(|e| Error::Candle(e))?;

        // Dequantize
        let dequantized_shifted = clamped
            .broadcast_sub(&zero_point_tensor)
            .map_err(|e| Error::Candle(e))?;
        let dequantized = dequantized_shifted
            .broadcast_mul(&scale_tensor)
            .map_err(|e| Error::Candle(e))?;

        Ok(dequantized)
    }

    /// Dynamic quantization (quantize weights only, activations stay FP32)
    fn quantize_tensor_dynamic(
        &self,
        tensor: &Tensor,
        scheme: &QuantizationScheme,
    ) -> Result<Tensor, Error> {
        // For dynamic quantization, we quantize and immediately dequantize
        self.quantize_tensor_int8(tensor, scheme)
    }

    /// Advanced quantization with external context support
    pub fn quantize_with_context(
        &self,
        model: &LoadedModel,
        context: &mut QuantizationContext,
        progress_callback: Option<ProgressFn>,
    ) -> Result<HashMap<String, Tensor>, Error> {
        let start_time = std::time::Instant::now();

        if let Some(ref callback) = progress_callback {
            callback(ProgressEvent::Status {
                message: "Starting advanced model quantization with context".to_string(),
            });
        }

        let mut quantized_tensors = HashMap::new();
        let total_tensors = model.raw_tensors.len();
        let mut total_original_size = 0usize;
        let mut total_quantized_size = 0usize;

        for (idx, (tensor_name, tensor)) in model.raw_tensors.iter().enumerate() {
            if let Some(ref callback) = progress_callback {
                if idx % 10 == 0 {
                    callback(ProgressEvent::Status {
                        message: format!(
                            "Quantizing tensor {}/{}: {}",
                            idx + 1,
                            total_tensors,
                            tensor_name
                        ),
                    });
                }
            }

            let quantized_tensor = if self.should_skip_layer(tensor_name) {
                tensor.clone()
            } else {
                // Check for context-specific scheme first
                let scheme = if let Some(context_scheme) = context.get_layer_scheme(tensor_name) {
                    context_scheme.clone()
                } else {
                    self.get_quantization_scheme(tensor_name, tensor)?
                };

                // Store tensor statistics in context if advanced stats enabled
                if self.config.advanced_stats {
                    let stats = self.collect_tensor_stats(tensor)?;
                    context.tensor_stats.insert(tensor_name.clone(), stats);
                }

                // Perform quantization
                if self.config.block_wise {
                    // For block-wise, store individual block schemes
                    let blocks = self.compute_tensor_blocks(tensor, &scheme)?;
                    context
                        .tensor_blocks
                        .insert(tensor_name.clone(), blocks.clone());
                    self.quantize_tensor_with_blocks(tensor, &blocks)?
                } else {
                    self.quantize_tensor(tensor, &scheme)?
                }
            };

            // Update metrics
            let original_size = tensor.elem_count() * 4; // Assume F32
            let quantized_size = match self.config.quantization_type {
                QuantizationType::Int8 => tensor.elem_count(),
                QuantizationType::Int4 => tensor.elem_count() / 2,
                _ => tensor.elem_count() * 4,
            };

            total_original_size += original_size;
            total_quantized_size += quantized_size;

            let compression_ratio = original_size as f32 / quantized_size as f32;
            context
                .metrics
                .layer_compression
                .insert(tensor_name.clone(), compression_ratio);

            quantized_tensors.insert(tensor_name.clone(), quantized_tensor);
        }

        // Update overall metrics
        context.metrics.compression_ratio =
            total_original_size as f32 / total_quantized_size as f32;
        context.metrics.quantization_time = start_time.elapsed().as_secs_f64();

        if let Some(ref callback) = progress_callback {
            callback(ProgressEvent::Status {
                message: format!(
                    "Advanced quantization completed. Compression ratio: {:.2}x",
                    context.metrics.compression_ratio
                ),
            });
        }

        Ok(quantized_tensors)
    }

    /// Compute block schemes for a tensor
    fn compute_tensor_blocks(
        &self,
        tensor: &Tensor,
        base_scheme: &QuantizationScheme,
    ) -> Result<Vec<QuantizationScheme>, Error> {
        let total_elements = tensor.shape().elem_count();
        let num_blocks = (total_elements + self.config.block_size - 1) / self.config.block_size;
        let mut blocks = Vec::new();

        let flat_tensor = tensor.flatten_all().map_err(|e| Error::Candle(e))?;
        let data: Vec<f32> = flat_tensor.to_vec1().map_err(|e| Error::Candle(e))?;

        for block_idx in 0..num_blocks {
            let start_idx = block_idx * self.config.block_size;
            let end_idx = (start_idx + self.config.block_size).min(data.len());
            let block_data = &data[start_idx..end_idx];

            let block_tensor =
                Tensor::new(block_data, tensor.device()).map_err(|e| Error::Candle(e))?;
            let block_scheme =
                self.compute_block_quantization_scheme(&block_tensor, base_scheme)?;
            blocks.push(block_scheme);
        }

        Ok(blocks)
    }

    /// Quantize tensor using pre-computed blocks
    fn quantize_tensor_with_blocks(
        &self,
        tensor: &Tensor,
        blocks: &[QuantizationScheme],
    ) -> Result<Tensor, Error> {
        let shape = tensor.shape().clone();
        let flat_tensor = tensor.flatten_all().map_err(|e| Error::Candle(e))?;
        let data: Vec<f32> = flat_tensor.to_vec1().map_err(|e| Error::Candle(e))?;

        let mut quantized_data = Vec::new();

        for (block_idx, block_scheme) in blocks.iter().enumerate() {
            let start_idx = block_idx * self.config.block_size;
            let end_idx = (start_idx + self.config.block_size).min(data.len());
            let block_data = &data[start_idx..end_idx];

            let block_tensor =
                Tensor::new(block_data, tensor.device()).map_err(|e| Error::Candle(e))?;

            let quantized_block = match block_scheme.quant_type {
                QuantizationType::Int8 => self.quantize_tensor_int8(&block_tensor, block_scheme)?,
                QuantizationType::Int4 => self.quantize_tensor_int4(&block_tensor, block_scheme)?,
                QuantizationType::Dynamic => {
                    self.quantize_tensor_dynamic(&block_tensor, block_scheme)?
                }
                _ => block_tensor,
            };

            let block_result: Vec<f32> = quantized_block.to_vec1().map_err(|e| Error::Candle(e))?;
            quantized_data.extend_from_slice(&block_result);
        }

        let result_tensor = Tensor::new(&quantized_data[..data.len()], tensor.device())
            .map_err(|e| Error::Candle(e))?
            .reshape(&shape)
            .map_err(|e| Error::Candle(e))?;

        Ok(result_tensor)
    }
}

/// Utilities for quantization-aware loading
pub mod quantized_loading {
    use super::*;

    /// Load a quantized model with proper handling
    pub fn load_quantized_model<P: AsRef<Path>>(
        path: P,
        options: LoadOptions,
        progress_callback: Option<ProgressFn>,
    ) -> Result<LoadedModel, Error> {
        // Load the model with options
        let model = crate::load_model(path.as_ref(), options)?;

        // TODO: Check if model has quantization metadata in config
        // For now, return the model as-is

        Ok(model)
    }

    /// Check if a model is quantized
    pub fn is_quantized(_model: &LoadedModel) -> bool {
        // TODO: Check model config for quantization metadata
        false
    }

    /// Get quantization information from model
    pub fn get_quantization_info(_model: &LoadedModel) -> Option<QuantizationConfig> {
        // TODO: Extract quantization config from model metadata
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_core::{DType, Device};

    #[test]
    fn test_quantization_config_default() {
        let config = QuantizationConfig::default();
        assert_eq!(config.quantization_type, QuantizationType::Int8);
        assert_eq!(config.calibration_samples, 128);
        assert_eq!(config.calibration_method, "minmax");
        assert!(config.symmetric);
    }

    #[test]
    fn test_quantization_scheme_default() {
        let scheme = QuantizationScheme::default();
        assert_eq!(scheme.quant_type, QuantizationType::Int8);
        assert_eq!(scheme.scale, 1.0);
        assert_eq!(scheme.zero_point, 0);
        assert!(scheme.symmetric);
        assert_eq!(scheme.range, (-128.0, 127.0));
    }

    #[test]
    fn test_calibration_dataset() {
        let dataset = CalibrationDataset::new(Device::Cpu);
        assert!(dataset.is_empty());
        assert_eq!(dataset.len(), 0);
    }

    #[test]
    fn test_quantization_engine_creation() {
        let config = QuantizationConfig::default();
        let engine = QuantizationEngine::new(config, Device::Cpu);
        assert_eq!(engine.config.quantization_type, QuantizationType::Int8);
    }
}