//! Advanced Quantization Methods
//!
//! This module implements state-of-the-art quantization techniques:
//! - AWQ (Activation-aware Weight Quantization): Finds optimal scales using activation importance
//! - GPTQ (Accurate Post-Training Quantization): Hessian-based blockwise quantization
//! - SmoothQuant: Migrates quantization difficulty from activations to weights
//!
//! These methods enable INT4/INT3 weight-only quantization with minimal accuracy loss.

use crate::{
    Result,
    error::{CompressionError, QuantizationError},
};
use rtx_tensor::{Device, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

// Helper to convert tensor errors
fn tensor_err(e: impl std::fmt::Display) -> CompressionError {
    CompressionError::Quantization(QuantizationError::TensorError(e.to_string()))
}

// =============================================================================
// AWQ (Activation-aware Weight Quantization)
// =============================================================================

/// Configuration for AWQ quantization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AWQConfig {
    /// Target bit width (typically 4 or 3)
    pub bit_width: u8,
    /// Group size for quantization (typically 128)
    pub group_size: usize,
    /// Number of calibration samples
    pub num_calibration_samples: usize,
    /// Percentile for activation importance (typically 0.9-0.99)
    pub activation_percentile: f32,
    /// Alpha for mixing search (0.0-1.0)
    pub alpha: f32,
    /// Number of search iterations for optimal scale
    pub num_search_iters: usize,
    /// Enable per-channel scaling
    pub per_channel: bool,
    /// Device for computation
    pub device: Device,
}

impl Default for AWQConfig {
    fn default() -> Self {
        Self {
            bit_width: 4,
            group_size: 128,
            num_calibration_samples: 128,
            activation_percentile: 0.95,
            alpha: 0.5,
            num_search_iters: 20,
            per_channel: true,
            device: Device::Cpu,
        }
    }
}

impl AWQConfig {
    /// Create new AWQ config with specified bit width
    pub fn new(bit_width: u8) -> Self {
        Self {
            bit_width,
            ..Default::default()
        }
    }

    /// Set group size
    pub fn with_group_size(mut self, group_size: usize) -> Self {
        self.group_size = group_size;
        self
    }

    /// Set calibration samples
    pub fn with_calibration_samples(mut self, num_samples: usize) -> Self {
        self.num_calibration_samples = num_samples;
        self
    }
}

/// AWQ Quantizer for activation-aware weight quantization
#[derive(Debug)]
pub struct AWQQuantizer {
    config: AWQConfig,
    /// Per-channel scales computed from activation importance
    activation_scales: HashMap<String, Vec<f32>>,
    /// Statistics collected during calibration
    stats: AWQStats,
}

/// Statistics from AWQ quantization
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AWQStats {
    /// Number of layers quantized
    pub layers_quantized: usize,
    /// Total weights quantized
    pub total_weights: usize,
    /// Average quantization error (MSE)
    pub avg_quant_error: f64,
    /// Salient channel ratios per layer
    pub salient_ratios: HashMap<String, f32>,
}

impl AWQQuantizer {
    /// Create new AWQ quantizer
    pub fn new(config: AWQConfig) -> Self {
        Self {
            config,
            activation_scales: HashMap::new(),
            stats: AWQStats::default(),
        }
    }

    /// Calibrate scales using activation data
    pub fn calibrate(&mut self, activations: &HashMap<String, Vec<Tensor>>) -> Result<()> {
        for (layer_name, act_samples) in activations {
            // Compute per-channel activation magnitudes
            let channel_magnitudes = self.compute_channel_magnitudes(act_samples)?;

            // Find salient channels (high activation magnitude)
            let percentile_threshold =
                self.compute_percentile(&channel_magnitudes, self.config.activation_percentile);

            // Compute scales based on activation importance
            let scales: Vec<f32> = channel_magnitudes
                .iter()
                .map(|&mag| {
                    if mag > percentile_threshold {
                        // Salient channel: use larger scale to preserve precision
                        (mag / percentile_threshold).sqrt()
                    } else {
                        1.0
                    }
                })
                .collect();

            self.activation_scales.insert(layer_name.clone(), scales);
        }
        Ok(())
    }

    /// Quantize weights using AWQ method
    pub fn quantize_weights(
        &mut self,
        layer_name: &str,
        weights: &Tensor,
    ) -> Result<AWQQuantizedWeight> {
        let shape = weights.shape().dims();
        let out_features = shape[0];
        let in_features = shape[1];

        // Get activation scales for this layer (or use uniform scales)
        let scales = self
            .activation_scales
            .get(layer_name)
            .cloned()
            .unwrap_or_else(|| vec![1.0; in_features]);

        // Scale weights by activation importance
        let mut weight_data = weights.to_vec().map_err(tensor_err)?;
        for j in 0..out_features {
            for i in 0..in_features {
                let idx = j * in_features + i;
                let scale = scales.get(i).copied().unwrap_or(1.0);
                weight_data[idx] *= scale;
            }
        }

        // Create scaled weights tensor for group quantization
        let scaled_weights = Tensor::from_slice(
            &weight_data,
            &[out_features, in_features],
            &self.config.device,
        )
        .map_err(tensor_err)?;

        // Quantize scaled weights using group quantization
        let (quantized, quant_scales, zeros) = self.group_quantize(&scaled_weights)?;

        // Compute inverse scales for dequantization
        let inv_activation_scales: Vec<f32> = scales.iter().map(|&s| 1.0 / s).collect();

        self.stats.layers_quantized += 1;
        self.stats.total_weights += out_features * in_features;

        Ok(AWQQuantizedWeight {
            quantized_data: quantized,
            scales: quant_scales,
            zeros,
            activation_scales: inv_activation_scales,
            group_size: self.config.group_size,
            bit_width: self.config.bit_width,
            shape: vec![out_features, in_features],
        })
    }

    /// Group-wise quantization of weights
    fn group_quantize(&self, weights: &Tensor) -> Result<(Vec<i8>, Vec<f32>, Vec<i8>)> {
        let shape = weights.shape().dims();
        let total_elements = shape.iter().product::<usize>();
        let weight_data = weights.to_vec().map_err(tensor_err)?;

        let mut quantized = Vec::with_capacity(total_elements);
        let mut scales = Vec::new();
        let mut zeros = Vec::new();

        let num_groups = (total_elements + self.config.group_size - 1) / self.config.group_size;

        for group_idx in 0..num_groups {
            let start = group_idx * self.config.group_size;
            let end = (start + self.config.group_size).min(total_elements);

            // Find min/max in this group
            let mut min_val = f32::MAX;
            let mut max_val = f32::MIN;
            for i in start..end {
                let val = weight_data[i];
                min_val = min_val.min(val);
                max_val = max_val.max(val);
            }

            // Compute scale and zero point
            let qmin = -(1 << (self.config.bit_width - 1));
            let qmax = (1 << (self.config.bit_width - 1)) - 1;

            let scale = (max_val - min_val) / (qmax - qmin) as f32;
            let scale = if scale == 0.0 { 1.0 } else { scale };
            let zero = ((qmin as f32 * scale - min_val) / scale).round() as i8;

            scales.push(scale);
            zeros.push(zero);

            // Quantize values in group
            for i in start..end {
                let val = weight_data[i];
                let q = ((val / scale) + zero as f32).round();
                let q = q.max(qmin as f32).min(qmax as f32) as i8;
                quantized.push(q);
            }
        }

        Ok((quantized, scales, zeros))
    }

    /// Compute per-channel activation magnitudes
    fn compute_channel_magnitudes(&self, activations: &[Tensor]) -> Result<Vec<f32>> {
        if activations.is_empty() {
            return Ok(Vec::new());
        }

        let shape = activations[0].shape().dims();
        let num_channels = if shape.len() >= 2 {
            shape[shape.len() - 1]
        } else {
            shape[0]
        };

        let mut magnitudes = vec![0.0f32; num_channels];
        let mut counts = vec![0usize; num_channels];

        for act in activations {
            let flat = act.to_vec().unwrap_or_default();
            for (i, &val) in flat.iter().enumerate() {
                let channel_idx = i % num_channels;
                magnitudes[channel_idx] += val.abs();
                counts[channel_idx] += 1;
            }
        }

        // Average magnitudes
        for i in 0..num_channels {
            if counts[i] > 0 {
                magnitudes[i] /= counts[i] as f32;
            }
        }

        Ok(magnitudes)
    }

    /// Compute percentile value
    fn compute_percentile(&self, values: &[f32], percentile: f32) -> f32 {
        if values.is_empty() {
            return 0.0;
        }
        let mut sorted = values.to_vec();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        let idx = ((sorted.len() as f32 - 1.0) * percentile) as usize;
        sorted[idx.min(sorted.len() - 1)]
    }

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

/// AWQ quantized weight representation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AWQQuantizedWeight {
    /// Quantized weight data (packed integers)
    pub quantized_data: Vec<i8>,
    /// Per-group scales
    pub scales: Vec<f32>,
    /// Per-group zero points
    pub zeros: Vec<i8>,
    /// Activation-derived inverse scales for dequant
    pub activation_scales: Vec<f32>,
    /// Group size
    pub group_size: usize,
    /// Bit width
    pub bit_width: u8,
    /// Original shape [out_features, in_features]
    pub shape: Vec<usize>,
}

impl AWQQuantizedWeight {
    /// Dequantize weights
    pub fn dequantize(&self) -> Result<Tensor> {
        let out_features = self.shape[0];
        let in_features = self.shape[1];
        let total = out_features * in_features;

        let mut dequantized = vec![0.0f32; total];
        let num_groups = self.scales.len();

        for (i, &qval) in self.quantized_data.iter().enumerate() {
            if i >= total {
                break;
            }

            let group_idx = i / self.group_size;
            let group_idx = group_idx.min(num_groups - 1);

            let scale = self.scales[group_idx];
            let zero = self.zeros[group_idx] as f32;
            let act_scale = self
                .activation_scales
                .get(i % in_features)
                .copied()
                .unwrap_or(1.0);

            dequantized[i] = (qval as f32 - zero) * scale * act_scale;
        }

        Tensor::from_slice(&dequantized, &[out_features, in_features], &Device::Cpu)
            .map_err(tensor_err)
    }

    /// Get memory footprint in bytes
    pub fn memory_bytes(&self) -> usize {
        let data_bytes = (self.quantized_data.len() * self.bit_width as usize + 7) / 8;
        let scale_bytes = self.scales.len() * 4; // f32
        let zero_bytes = self.zeros.len(); // i8
        let act_scale_bytes = self.activation_scales.len() * 4; // f32
        data_bytes + scale_bytes + zero_bytes + act_scale_bytes
    }
}

// =============================================================================
// GPTQ (Accurate Post-Training Quantization)
// =============================================================================

/// Configuration for GPTQ quantization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GPTQConfig {
    /// Target bit width (typically 4 or 3)
    pub bit_width: u8,
    /// Group size for quantization
    pub group_size: usize,
    /// Block size for Hessian computation
    pub block_size: usize,
    /// Damping factor for Hessian (prevents numerical instability)
    pub damp_percent: f32,
    /// Enable static groups (same group across different rows)
    pub static_groups: bool,
    /// Number of calibration samples
    pub num_calibration_samples: usize,
    /// Enable activation order optimization
    pub act_order: bool,
    /// Use true sequential (quantize one column at a time)
    pub true_sequential: bool,
    /// Device for computation
    pub device: Device,
}

impl Default for GPTQConfig {
    fn default() -> Self {
        Self {
            bit_width: 4,
            group_size: 128,
            block_size: 128,
            damp_percent: 0.01,
            static_groups: true,
            num_calibration_samples: 128,
            act_order: false,
            true_sequential: true,
            device: Device::Cpu,
        }
    }
}

impl GPTQConfig {
    /// Create new GPTQ config
    pub fn new(bit_width: u8) -> Self {
        Self {
            bit_width,
            ..Default::default()
        }
    }

    /// Set group size
    pub fn with_group_size(mut self, group_size: usize) -> Self {
        self.group_size = group_size;
        self
    }

    /// Enable activation order optimization
    pub fn with_act_order(mut self, enabled: bool) -> Self {
        self.act_order = enabled;
        self
    }
}

/// GPTQ Quantizer using Hessian-based optimization
#[derive(Debug)]
pub struct GPTQQuantizer {
    config: GPTQConfig,
    /// Cached Hessian matrices per layer
    hessians: HashMap<String, Vec<Vec<f32>>>,
    /// Statistics
    stats: GPTQStats,
}

/// Statistics from GPTQ quantization
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct GPTQStats {
    /// Number of layers quantized
    pub layers_quantized: usize,
    /// Total weights quantized
    pub total_weights: usize,
    /// Average quantization error before compensation
    pub avg_error_before: f64,
    /// Average quantization error after compensation
    pub avg_error_after: f64,
    /// Hessian computation time (ms)
    pub hessian_time_ms: f64,
    /// Quantization time (ms)
    pub quant_time_ms: f64,
}

impl GPTQQuantizer {
    /// Create new GPTQ quantizer
    pub fn new(config: GPTQConfig) -> Self {
        Self {
            config,
            hessians: HashMap::new(),
            stats: GPTQStats::default(),
        }
    }

    /// Calibrate Hessian using input activations
    pub fn calibrate(&mut self, layer_name: &str, inputs: &[Tensor]) -> Result<()> {
        if inputs.is_empty() {
            return Ok(());
        }

        let in_features = inputs[0].shape().dims().last().copied().unwrap_or(1);

        // Initialize Hessian as H = X^T * X (accumulated)
        let mut hessian = vec![vec![0.0f32; in_features]; in_features];

        for input in inputs {
            let flat = input.to_vec().map_err(tensor_err)?;
            let batch_size = flat.len() / in_features;

            for b in 0..batch_size {
                for i in 0..in_features {
                    for j in 0..in_features {
                        let xi = flat[b * in_features + i];
                        let xj = flat[b * in_features + j];
                        hessian[i][j] += xi * xj;
                    }
                }
            }
        }

        // Normalize by number of samples
        let num_samples = inputs.iter().map(|t| t.shape().dims()[0]).sum::<usize>();
        for i in 0..in_features {
            for j in 0..in_features {
                hessian[i][j] /= num_samples as f32;
            }
        }

        // Add damping to diagonal
        let diag_mean: f32 =
            (0..in_features).map(|i| hessian[i][i]).sum::<f32>() / in_features as f32;
        let damp = self.config.damp_percent * diag_mean;
        for i in 0..in_features {
            hessian[i][i] += damp;
        }

        self.hessians.insert(layer_name.to_string(), hessian);
        Ok(())
    }

    /// Quantize weights using GPTQ algorithm
    pub fn quantize_weights(
        &mut self,
        layer_name: &str,
        weights: &Tensor,
    ) -> Result<GPTQQuantizedWeight> {
        let shape = weights.shape().dims();
        let out_features = shape[0];
        let in_features = shape[1];

        // Get Hessian (or create identity if not calibrated)
        let hessian = self.hessians.get(layer_name).cloned().unwrap_or_else(|| {
            let mut h = vec![vec![0.0f32; in_features]; in_features];
            for i in 0..in_features {
                h[i][i] = 1.0;
            }
            h
        });

        // Compute Cholesky decomposition of Hessian inverse
        let h_inv = self.cholesky_inverse(&hessian)?;

        // Working copy of weights
        let mut w = weights.to_vec().map_err(tensor_err)?;
        let mut quantized = vec![0i8; out_features * in_features];
        let mut scales = Vec::new();
        let mut zeros = Vec::new();

        // Determine column order (optionally by activation importance)
        let col_order: Vec<usize> = if self.config.act_order {
            // Sort by diagonal of Hessian (activation importance)
            let mut indexed: Vec<_> = (0..in_features).map(|i| (i, hessian[i][i])).collect();
            indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
            indexed.into_iter().map(|(i, _)| i).collect()
        } else {
            (0..in_features).collect()
        };

        // Quantize column by column
        for (block_start, _) in (0..in_features).step_by(self.config.block_size).enumerate() {
            let block_end = (block_start + self.config.block_size).min(in_features);

            for col_idx in block_start..block_end {
                let col = col_order[col_idx];

                // Compute group index
                let group_idx = col / self.config.group_size;

                // Find min/max for this column across all rows
                let mut col_min = f32::MAX;
                let mut col_max = f32::MIN;
                for row in 0..out_features {
                    let val = w[row * in_features + col];
                    col_min = col_min.min(val);
                    col_max = col_max.max(val);
                }

                // Compute scale and zero
                let qmin = -(1 << (self.config.bit_width - 1));
                let qmax = (1 << (self.config.bit_width - 1)) - 1;
                let scale = (col_max - col_min) / (qmax - qmin) as f32;
                let scale = if scale == 0.0 { 1.0 } else { scale };
                let zero = ((qmin as f32 * scale - col_min) / scale).round() as i8;

                if group_idx >= scales.len() {
                    scales.push(scale);
                    zeros.push(zero);
                }

                // Quantize this column for each row
                for row in 0..out_features {
                    let idx = row * in_features + col;
                    let val = w[idx];
                    let q = ((val / scale) + zero as f32).round();
                    let q = q.max(qmin as f32).min(qmax as f32);
                    quantized[idx] = q as i8;

                    // Compute quantization error
                    let dequant = (q - zero as f32) * scale;
                    let error = val - dequant;

                    // Update remaining weights to compensate for error (GPTQ key insight)
                    // w_remaining -= error * H_inv[col, remaining] / H_inv[col, col]
                    let h_diag = h_inv[col][col];
                    if h_diag > 1e-10 {
                        for remaining_col in (col_idx + 1)..in_features {
                            let rem = col_order[remaining_col];
                            let update = error * h_inv[col][rem] / h_diag;
                            w[row * in_features + rem] -= update;
                        }
                    }
                }
            }
        }

        self.stats.layers_quantized += 1;
        self.stats.total_weights += out_features * in_features;

        Ok(GPTQQuantizedWeight {
            quantized_data: quantized,
            scales,
            zeros,
            group_size: self.config.group_size,
            bit_width: self.config.bit_width,
            shape: vec![out_features, in_features],
            col_order: if self.config.act_order {
                Some(col_order)
            } else {
                None
            },
        })
    }

    /// Compute Cholesky-based inverse of a symmetric positive definite matrix
    fn cholesky_inverse(&self, matrix: &[Vec<f32>]) -> Result<Vec<Vec<f32>>> {
        let n = matrix.len();

        // Simple pseudo-inverse using regularized SVD-like approach
        // For production, use proper linear algebra library
        let mut inv = vec![vec![0.0f32; n]; n];

        for i in 0..n {
            let diag = matrix[i][i];
            if diag > 1e-10 {
                inv[i][i] = 1.0 / diag;
            } else {
                inv[i][i] = 1.0;
            }
        }

        // For off-diagonal elements, use approximate inverse
        for i in 0..n {
            for j in 0..n {
                if i != j && matrix[i][i] > 1e-10 && matrix[j][j] > 1e-10 {
                    inv[i][j] = -matrix[i][j] / (matrix[i][i] * matrix[j][j]).sqrt();
                }
            }
        }

        Ok(inv)
    }

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

/// GPTQ quantized weight representation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GPTQQuantizedWeight {
    /// Quantized weight data
    pub quantized_data: Vec<i8>,
    /// Per-group scales
    pub scales: Vec<f32>,
    /// Per-group zero points
    pub zeros: Vec<i8>,
    /// Group size
    pub group_size: usize,
    /// Bit width
    pub bit_width: u8,
    /// Original shape
    pub shape: Vec<usize>,
    /// Column order (if act_order enabled)
    pub col_order: Option<Vec<usize>>,
}

impl GPTQQuantizedWeight {
    /// Dequantize weights
    pub fn dequantize(&self) -> Result<Tensor> {
        let out_features = self.shape[0];
        let in_features = self.shape[1];

        let mut dequantized = vec![0.0f32; out_features * in_features];

        for row in 0..out_features {
            for col in 0..in_features {
                let idx = row * in_features + col;
                let group_idx = col / self.group_size;
                let group_idx = group_idx.min(self.scales.len() - 1);

                let q = self.quantized_data[idx] as f32;
                let scale = self.scales[group_idx];
                let zero = self.zeros[group_idx] as f32;

                dequantized[idx] = (q - zero) * scale;
            }
        }

        Tensor::from_slice(&dequantized, &[out_features, in_features], &Device::Cpu)
            .map_err(tensor_err)
    }

    /// Get compression ratio
    pub fn compression_ratio(&self) -> f32 {
        let original_bits = self.shape[0] * self.shape[1] * 32; // FP32
        let quantized_bits = self.quantized_data.len() * self.bit_width as usize
            + self.scales.len() * 32
            + self.zeros.len() * 8;
        original_bits as f32 / quantized_bits as f32
    }
}

// =============================================================================
// SmoothQuant
// =============================================================================

/// Configuration for SmoothQuant
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SmoothQuantConfig {
    /// Smoothing alpha (0.0-1.0, higher = more smoothing on activations)
    pub alpha: f32,
    /// Target bit width for weights
    pub weight_bit_width: u8,
    /// Target bit width for activations
    pub activation_bit_width: u8,
    /// Number of calibration samples
    pub num_calibration_samples: usize,
    /// Per-channel smoothing
    pub per_channel: bool,
    /// Minimum scale to prevent division by zero
    pub min_scale: f32,
    /// Device for computation
    pub device: Device,
}

impl Default for SmoothQuantConfig {
    fn default() -> Self {
        Self {
            alpha: 0.5,
            weight_bit_width: 8,
            activation_bit_width: 8,
            num_calibration_samples: 128,
            per_channel: true,
            min_scale: 1e-5,
            device: Device::Cpu,
        }
    }
}

impl SmoothQuantConfig {
    /// Create new SmoothQuant config
    pub fn new(alpha: f32) -> Self {
        Self {
            alpha: alpha.clamp(0.0, 1.0),
            ..Default::default()
        }
    }

    /// Set bit widths
    pub fn with_bit_widths(mut self, weight_bits: u8, activation_bits: u8) -> Self {
        self.weight_bit_width = weight_bits;
        self.activation_bit_width = activation_bits;
        self
    }
}

/// SmoothQuant migrates quantization difficulty from activations to weights
#[derive(Debug)]
pub struct SmoothQuantQuantizer {
    config: SmoothQuantConfig,
    /// Smoothing scales per layer (applied to activations)
    smoothing_scales: HashMap<String, Vec<f32>>,
    /// Statistics
    stats: SmoothQuantStats,
}

/// Statistics from SmoothQuant
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct SmoothQuantStats {
    /// Number of layers smoothed
    pub layers_smoothed: usize,
    /// Average activation range before smoothing
    pub avg_act_range_before: f64,
    /// Average activation range after smoothing
    pub avg_act_range_after: f64,
    /// Average weight range before smoothing
    pub avg_weight_range_before: f64,
    /// Average weight range after smoothing
    pub avg_weight_range_after: f64,
}

impl SmoothQuantQuantizer {
    /// Create new SmoothQuant quantizer
    pub fn new(config: SmoothQuantConfig) -> Self {
        Self {
            config,
            smoothing_scales: HashMap::new(),
            stats: SmoothQuantStats::default(),
        }
    }

    /// Calibrate smoothing scales using activations and weights
    pub fn calibrate(
        &mut self,
        layer_name: &str,
        activations: &[Tensor],
        weights: &Tensor,
    ) -> Result<()> {
        if activations.is_empty() {
            return Ok(());
        }

        let in_features = weights.shape().dims()[1];

        // Compute per-channel activation max
        let mut act_max = vec![0.0f32; in_features];
        for act in activations {
            let flat = act.to_vec().map_err(tensor_err)?;
            let num_elements = flat.len();
            for (i, &val) in flat.iter().enumerate() {
                let channel = i % in_features;
                act_max[channel] = act_max[channel].max(val.abs());
            }
        }

        // Compute per-channel weight max (for each input channel)
        let weight_flat = weights.to_vec().map_err(tensor_err)?;
        let out_features = weights.shape().dims()[0];
        let mut weight_max = vec![0.0f32; in_features];
        for out_idx in 0..out_features {
            for in_idx in 0..in_features {
                let val = weight_flat[out_idx * in_features + in_idx];
                weight_max[in_idx] = weight_max[in_idx].max(val.abs());
            }
        }

        // Compute smoothing scales: s = act_max^alpha / weight_max^(1-alpha)
        let alpha = self.config.alpha;
        let scales: Vec<f32> = (0..in_features)
            .map(|i| {
                let a = act_max[i].max(self.config.min_scale);
                let w = weight_max[i].max(self.config.min_scale);
                (a.powf(alpha) / w.powf(1.0 - alpha)).max(self.config.min_scale)
            })
            .collect();

        self.smoothing_scales.insert(layer_name.to_string(), scales);
        self.stats.layers_smoothed += 1;

        Ok(())
    }

    /// Get smoothing scales for a layer
    pub fn get_scales(&self, layer_name: &str) -> Option<&Vec<f32>> {
        self.smoothing_scales.get(layer_name)
    }

    /// Apply smoothing to activations (multiply by inverse scale)
    pub fn smooth_activations(&self, layer_name: &str, activations: &Tensor) -> Result<Tensor> {
        let scales = self.smoothing_scales.get(layer_name).ok_or_else(|| {
            CompressionError::Quantization(QuantizationError::TensorError(format!(
                "No scales for layer {}",
                layer_name
            )))
        })?;

        let shape = activations.shape().dims();
        let in_features = shape[shape.len() - 1];

        let mut smoothed = activations.to_vec().map_err(tensor_err)?;
        for (i, val) in smoothed.iter_mut().enumerate() {
            let channel = i % in_features;
            let inv_scale = 1.0 / scales.get(channel).copied().unwrap_or(1.0);
            *val *= inv_scale;
        }

        Tensor::from_slice(&smoothed, shape, &self.config.device).map_err(tensor_err)
    }

    /// Apply inverse smoothing to weights (multiply by scale)
    pub fn smooth_weights(&self, layer_name: &str, weights: &Tensor) -> Result<Tensor> {
        let scales = self.smoothing_scales.get(layer_name).ok_or_else(|| {
            CompressionError::Quantization(QuantizationError::TensorError(format!(
                "No scales for layer {}",
                layer_name
            )))
        })?;

        let shape = weights.shape().dims();
        let out_features = shape[0];
        let in_features = shape[1];

        let mut smoothed = weights.to_vec().map_err(tensor_err)?;
        for out_idx in 0..out_features {
            for in_idx in 0..in_features {
                let idx = out_idx * in_features + in_idx;
                let scale = scales.get(in_idx).copied().unwrap_or(1.0);
                smoothed[idx] *= scale;
            }
        }

        Tensor::from_slice(&smoothed, shape, &self.config.device).map_err(tensor_err)
    }

    /// Quantize a smoothed layer (both weights and activations)
    pub fn quantize_layer(
        &mut self,
        layer_name: &str,
        weights: &Tensor,
        activations: &[Tensor],
    ) -> Result<SmoothQuantizedLayer> {
        // First calibrate
        self.calibrate(layer_name, activations, weights)?;

        // Smooth weights
        let smoothed_weights = self.smooth_weights(layer_name, weights)?;

        // Quantize smoothed weights
        let quantized_weights =
            self.quantize_tensor(&smoothed_weights, self.config.weight_bit_width)?;

        let scales = self
            .smoothing_scales
            .get(layer_name)
            .cloned()
            .unwrap_or_default();

        Ok(SmoothQuantizedLayer {
            quantized_weights,
            smoothing_scales: scales,
            weight_bit_width: self.config.weight_bit_width,
            activation_bit_width: self.config.activation_bit_width,
        })
    }

    /// Simple tensor quantization
    fn quantize_tensor(&self, tensor: &Tensor, bit_width: u8) -> Result<QuantizedTensorData> {
        let flat = tensor.to_vec().map_err(tensor_err)?;

        // Find min/max
        let min_val = flat.iter().copied().fold(f32::MAX, f32::min);
        let max_val = flat.iter().copied().fold(f32::MIN, f32::max);

        let qmin = -(1i32 << (bit_width - 1));
        let qmax = (1i32 << (bit_width - 1)) - 1;

        let scale = (max_val - min_val) / (qmax - qmin) as f32;
        let scale = if scale == 0.0 { 1.0 } else { scale };
        let zero = ((qmin as f32 * scale - min_val) / scale).round() as i8;

        let quantized: Vec<i8> = flat
            .iter()
            .map(|&v| {
                let q = ((v / scale) + zero as f32).round();
                q.max(qmin as f32).min(qmax as f32) as i8
            })
            .collect();

        Ok(QuantizedTensorData {
            data: quantized,
            scale,
            zero_point: zero,
            shape: tensor.shape().dims().to_vec(),
        })
    }

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

/// Simple quantized tensor data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuantizedTensorData {
    /// Quantized data
    pub data: Vec<i8>,
    /// Scale factor
    pub scale: f32,
    /// Zero point
    pub zero_point: i8,
    /// Original shape
    pub shape: Vec<usize>,
}

impl QuantizedTensorData {
    /// Dequantize
    pub fn dequantize(&self, device: Device) -> Result<Tensor> {
        let dequantized: Vec<f32> = self
            .data
            .iter()
            .map(|&q| (q as f32 - self.zero_point as f32) * self.scale)
            .collect();

        Tensor::from_slice(&dequantized, &self.shape, &device)
            .map_err(|e| QuantizationError::TensorError(e.to_string()).into())
    }
}

/// SmoothQuant quantized layer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SmoothQuantizedLayer {
    /// Quantized weights
    pub quantized_weights: QuantizedTensorData,
    /// Smoothing scales for activation preprocessing
    pub smoothing_scales: Vec<f32>,
    /// Weight bit width
    pub weight_bit_width: u8,
    /// Activation bit width
    pub activation_bit_width: u8,
}

impl SmoothQuantizedLayer {
    /// Get inverse smoothing scales (for activation preprocessing)
    pub fn inverse_scales(&self) -> Vec<f32> {
        self.smoothing_scales.iter().map(|&s| 1.0 / s).collect()
    }

    /// Number of output features (rows of the weight matrix).
    #[inline]
    pub fn out_features(&self) -> usize {
        self.quantized_weights.shape[0]
    }

    /// Number of input features (columns of the weight matrix).
    #[inline]
    pub fn in_features(&self) -> usize {
        self.quantized_weights.shape[1]
    }

    /// Run the quantized linear layer forward pass on a raw f32 buffer.
    ///
    /// Implements the SmoothQuant inference path:
    /// 1. **Smooth activations** – divide each channel by its smoothing scale, migrating
    ///    quantisation difficulty from activations to the pre-scaled weights.
    /// 2. **Quantise activations to INT8** – compute a symmetric per-tensor scale
    ///    `act_scale = max(|smoothed|) / 127` and round to the range `[−127, 127]`.
    /// 3. **INT8 GEMM** – compute `acc[b, out] = Σ_k act_q[b,k] * (weight_q[out,k] − zp)`
    ///    with i32 accumulation.
    /// 4. **Dequantise** – `output[b, out] = acc[b, out] as f32 * act_scale * weight_scale`.
    ///
    /// # INT8 range
    ///
    /// Activations are clamped to **`−127…127`** (symmetric), matching the
    /// `activation_bit_width = 8` convention used throughout this crate.
    ///
    /// # Dequantisation formula
    ///
    /// ```text
    /// output[b, out] = (Σ_k act_q[b,k] * (weight_q[out,k] − zp)) * act_scale * weight_scale
    /// ```
    ///
    /// # Arguments
    ///
    /// * `activations` – Row-major f32 slice of shape `[batch_size, in_features]`.
    /// * `batch_size`  – Number of input rows.
    ///
    /// # Returns
    ///
    /// Row-major `Vec<f32>` of shape `[batch_size, out_features]`.
    ///
    /// # Errors
    ///
    /// Returns an error if the activation buffer length is inconsistent with `batch_size`.
    pub fn forward_raw(&self, activations: &[f32], batch_size: usize) -> Result<Vec<f32>> {
        let in_f = self.in_features();
        let out_f = self.out_features();

        let expected_len = batch_size * in_f;
        if activations.len() != expected_len {
            return Err(CompressionError::Quantization(
                QuantizationError::InvalidConfig(format!(
                    "SmoothQuantizedLayer::forward_raw: activation buffer length {} \
                     does not match batch_size={} × in_features={}",
                    activations.len(),
                    batch_size,
                    in_f
                )),
            ));
        }

        let zp = self.quantized_weights.zero_point as i32;
        let weight_scale = self.quantized_weights.scale;

        // ------------------------------------------------------------------
        // Step 1 + 2: smooth then quantise activations to INT8
        //
        // We compute smoothed values first (in place into a temporary buffer),
        // then find max_abs and derive act_scale before quantising.
        // ------------------------------------------------------------------
        let mut smoothed = vec![0.0f32; batch_size * in_f];
        for b in 0..batch_size {
            let row = b * in_f;
            for c in 0..in_f {
                // Divide by the per-channel smoothing scale (≥ ε for numerical safety).
                let s = self.smoothing_scales[c].max(f32::EPSILON);
                smoothed[row + c] = activations[row + c] / s;
            }
        }

        // Per-tensor activation scale: max absolute value over the whole batch.
        let max_abs = smoothed.iter().map(|v| v.abs()).fold(0.0f32, f32::max);

        // Guard against zero-tensor inputs; any non-zero scale works here.
        let act_scale = if max_abs < f32::EPSILON {
            1.0f32
        } else {
            max_abs / 127.0f32
        };

        // Quantise each smoothed value to INT8 in range [−127, 127].
        let act_q: Vec<i8> = smoothed
            .iter()
            .map(|&v| {
                let q = (v / act_scale).round();
                q.clamp(-127.0, 127.0) as i8
            })
            .collect();

        // ------------------------------------------------------------------
        // Step 3: INT8 GEMM with i32 accumulation
        //
        // output_i32[b, out_row] = Σ_k act_q[b,k] * (weight_q[out_row,k] − zp)
        // ------------------------------------------------------------------
        let mut output = vec![0.0f32; batch_size * out_f];
        for b in 0..batch_size {
            let act_row = b * in_f;
            for out_row in 0..out_f {
                let weight_row = out_row * in_f;
                let mut acc = 0i32;
                for k in 0..in_f {
                    let wq = self.quantized_weights.data[weight_row + k] as i32 - zp;
                    let aq = act_q[act_row + k] as i32;
                    acc += aq * wq;
                }
                // Step 4: dequantise
                output[b * out_f + out_row] = acc as f32 * act_scale * weight_scale;
            }
        }

        Ok(output)
    }
}

// =============================================================================
// Tests
// =============================================================================

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

    #[test]
    fn test_awq_config_default() {
        let config = AWQConfig::default();
        assert_eq!(config.bit_width, 4);
        assert_eq!(config.group_size, 128);
    }

    #[test]
    fn test_awq_quantizer_basic() {
        let config = AWQConfig::new(4).with_group_size(64);
        let quantizer = AWQQuantizer::new(config);
        assert_eq!(quantizer.stats().layers_quantized, 0);
    }

    #[test]
    fn test_gptq_config_default() {
        let config = GPTQConfig::default();
        assert_eq!(config.bit_width, 4);
        assert_eq!(config.group_size, 128);
        assert!(config.true_sequential);
    }

    #[test]
    fn test_gptq_quantizer_basic() {
        let config = GPTQConfig::new(4).with_act_order(true);
        let quantizer = GPTQQuantizer::new(config);
        assert!(quantizer.config.act_order);
    }

    #[test]
    fn test_smoothquant_config_default() {
        let config = SmoothQuantConfig::default();
        assert_eq!(config.alpha, 0.5);
        assert_eq!(config.weight_bit_width, 8);
        assert_eq!(config.activation_bit_width, 8);
    }

    #[test]
    fn test_smoothquant_alpha_clamping() {
        let config1 = SmoothQuantConfig::new(1.5);
        assert_eq!(config1.alpha, 1.0);

        let config2 = SmoothQuantConfig::new(-0.5);
        assert_eq!(config2.alpha, 0.0);
    }

    #[test]
    fn test_quantized_tensor_dequantize() {
        let data = QuantizedTensorData {
            data: vec![0, 10, -10, 5],
            scale: 0.1,
            zero_point: 0,
            shape: vec![2, 2],
        };

        let dequant = data.dequantize(Device::Cpu).unwrap();
        let flat: Vec<f32> = dequant.to_vec().unwrap();

        assert!((flat[0] - 0.0).abs() < 0.01);
        assert!((flat[1] - 1.0).abs() < 0.01);
        assert!((flat[2] - -1.0).abs() < 0.01);
        assert!((flat[3] - 0.5).abs() < 0.01);
    }

    #[test]
    fn test_awq_quantized_weight_memory() {
        let weight = AWQQuantizedWeight {
            quantized_data: vec![0; 1024],
            scales: vec![1.0; 8],
            zeros: vec![0; 8],
            activation_scales: vec![1.0; 256],
            group_size: 128,
            bit_width: 4,
            shape: vec![4, 256],
        };

        let bytes = weight.memory_bytes();
        assert!(bytes > 0);
        assert!(bytes < 1024 * 4 + 100); // Less than full FP32 + overhead
    }

    #[test]
    fn test_gptq_compression_ratio() {
        let weight = GPTQQuantizedWeight {
            quantized_data: vec![0; 1024],
            scales: vec![1.0; 8],
            zeros: vec![0; 8],
            group_size: 128,
            bit_width: 4,
            shape: vec![32, 32],
            col_order: None,
        };

        let ratio = weight.compression_ratio();
        assert!(ratio > 1.0); // Should be compressed
        assert!(ratio < 10.0); // But not impossibly so
    }

    // -----------------------------------------------------------------------
    // Helper: build a SmoothQuantizedLayer with fully controlled state
    //
    // Parameters:
    //   out_features, in_features — weight matrix shape
    //   weight_data   — INT8 quantized weights, row-major [out, in]
    //   weight_scale  — single global scale for the weight tensor
    //   weight_zp     — zero point for weight dequantization
    //   smoothing     — per-channel smoothing scales, len = in_features
    // -----------------------------------------------------------------------
    fn make_smooth_layer(
        out_features: usize,
        in_features: usize,
        weight_data: Vec<i8>,
        weight_scale: f32,
        weight_zp: i8,
        smoothing: Vec<f32>,
    ) -> SmoothQuantizedLayer {
        SmoothQuantizedLayer {
            quantized_weights: QuantizedTensorData {
                data: weight_data,
                scale: weight_scale,
                zero_point: weight_zp,
                shape: vec![out_features, in_features],
            },
            smoothing_scales: smoothing,
            weight_bit_width: 8,
            activation_bit_width: 8,
        }
    }

    /// Output length must equal batch_size × out_features.
    #[test]
    fn test_smoothquant_forward_raw_shape() {
        let out_f = 3usize;
        let in_f = 4usize;
        let layer = make_smooth_layer(
            out_f,
            in_f,
            vec![0i8; out_f * in_f],
            1.0,
            0,
            vec![1.0f32; in_f],
        );
        let activations = vec![0.0f32; 2 * in_f]; // batch=2
        let out = layer.forward_raw(&activations, 2).unwrap();
        assert_eq!(
            out.len(),
            2 * out_f,
            "output length must be batch * out_features"
        );
    }

    /// Identity layer (W=I, smoothing=1, zp=0, weight_scale=1) reproduces activations.
    ///
    /// With weight_scale=1 and act_scale derived from the input, roundtrip quantisation
    /// introduces < 1% relative error for values in [1, 127].
    #[test]
    fn test_smoothquant_forward_raw_identity_layer() {
        // 4×4 identity weight matrix (INT8), zero-point=0, weight_scale=1.
        // smoothing_scales = [1.0; 4] → no smoothing effect.
        // Activations: [1, 2, 3, 4] (batch=1).
        // Expected output ≈ [1, 2, 3, 4] (within quantisation tolerance).
        let in_f = 4usize;
        let out_f = 4usize;
        let identity: Vec<i8> = vec![
            1, 0, 0, 0, //
            0, 1, 0, 0, //
            0, 0, 1, 0, //
            0, 0, 0, 1, //
        ];
        let layer = make_smooth_layer(out_f, in_f, identity, 1.0, 0, vec![1.0f32; in_f]);
        let activations = vec![1.0f32, 2.0, 3.0, 4.0];
        let out = layer.forward_raw(&activations, 1).unwrap();

        assert_eq!(out.len(), out_f);
        // Tolerance: INT8 quantisation of activations loses at most 1/127 ≈ 0.8%.
        for (i, (&computed, &expected)) in out.iter().zip(activations.iter()).enumerate() {
            let rel_err = (computed - expected).abs() / (expected.abs().max(1e-6_f32));
            assert!(
                rel_err < 0.02,
                "element {i}: got {computed:.4}, expected {expected:.4}, rel_err={rel_err:.4}"
            );
        }
    }

    /// smoothing_scales=[2.0]*n divides each activation channel by 2 before quantisation.
    ///
    /// With W=I, this halves the output compared to smoothing_scales=[1.0].
    #[test]
    fn test_smoothquant_forward_raw_scale_effect() {
        let in_f = 4usize;
        let out_f = 4usize;
        let identity: Vec<i8> = vec![
            1, 0, 0, 0, //
            0, 1, 0, 0, //
            0, 0, 1, 0, //
            0, 0, 0, 1, //
        ];

        // Layer with smoothing_scales = [1.0; 4] (baseline)
        let layer_no_smooth =
            make_smooth_layer(out_f, in_f, identity.clone(), 1.0, 0, vec![1.0f32; in_f]);
        // Layer with smoothing_scales = [2.0; 4] (halves activations)
        let layer_smooth = make_smooth_layer(out_f, in_f, identity, 1.0, 0, vec![2.0f32; in_f]);

        let activations = vec![10.0f32, 20.0, 30.0, 40.0];

        let out_baseline = layer_no_smooth.forward_raw(&activations, 1).unwrap();
        let out_halved = layer_smooth.forward_raw(&activations, 1).unwrap();

        for (i, (&baseline, &halved)) in out_baseline.iter().zip(out_halved.iter()).enumerate() {
            let ratio = baseline / halved;
            assert!(
                (ratio - 2.0).abs() < 0.05,
                "element {i}: baseline={baseline:.4} / halved={halved:.4} = ratio {ratio:.4}, expected ~2.0"
            );
        }
    }

    /// Manually computed reference: 2×2 layer, known weights/scales/input.
    ///
    /// Setup:
    ///   W (INT8) = [[2, 0], [0, 2]], weight_scale=0.5, weight_zp=0
    ///   smoothing_scales = [1.0, 1.0]
    ///   activations = [10.0, 20.0] (batch=1)
    ///
    /// Step-by-step:
    ///   1. Smooth: [10/1, 20/1] = [10.0, 20.0]
    ///   2. act_scale = max(|10|, |20|) / 127 = 20/127 ≈ 0.15748
    ///   3. Quantise: act_q = round([10/0.15748, 20/0.15748]).clamp(-127,127)
    ///                      ≈ round([63.5, 127.0]) = [64, 127]  (or [63, 127] depending on rounding)
    ///   4. INT8 GEMM (weight_zp=0):
    ///      out[0] = 2*act_q[0] + 0*act_q[1]  = 2*64 = 128  (or 2*63=126)
    ///      out[1] = 0*act_q[0] + 2*act_q[1]  = 2*127 = 254
    ///   5. Dequant: output * act_scale * weight_scale
    ///      out[0] ≈ 128 * 0.15748 * 0.5 ≈ 10.08  (expected ~10.0)
    ///      out[1] ≈ 254 * 0.15748 * 0.5 ≈ 20.0
    ///
    /// Tolerance: 2% (quantisation rounding).
    #[test]
    fn test_smoothquant_forward_raw_matches_manual() {
        let layer = make_smooth_layer(2, 2, vec![2i8, 0, 0, 2], 0.5, 0, vec![1.0f32, 1.0]);
        let activations = vec![10.0f32, 20.0];
        let out = layer.forward_raw(&activations, 1).unwrap();

        assert_eq!(out.len(), 2);
        assert!(
            (out[0] - 10.0).abs() < 0.5,
            "out[0] ≈ 10.0, got {:.4}",
            out[0]
        );
        assert!(
            (out[1] - 20.0).abs() < 0.5,
            "out[1] ≈ 20.0, got {:.4}",
            out[1]
        );
    }
}