//! Quantization module for model compression.

use forge_shared::{ModelInfo, QuantMethod, QuantPrecision, QuantizationConfig};

/// Quantizer for model weight quantization.
#[derive(Debug)]
pub struct Quantizer {
    /// RNG state for calibration.
    rng_state: u64,
}

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

impl Quantizer {
    /// Create a new quantizer.
    pub fn new() -> Self {
        Self { rng_state: 42 }
    }

    /// Quantize a model and return compression factor.
    pub fn quantize(&mut self, model: &ModelInfo, config: &QuantizationConfig) -> f32 {
        // Calculate bit reduction
        let original_bits = model.precision_bits as f32;
        let target_bits = self.precision_to_bits(config.precision);

        // Base compression from bit reduction
        let mut compression = original_bits / target_bits;

        // Adjust for quantization method overhead
        compression *= self.method_efficiency(config.method);

        // Adjust for group quantization (adds scale/zero-point storage)
        if let Some(group_size) = config.group_size {
            let group_overhead = 1.0 + (2.0 * 16.0) / (group_size as f32 * target_bits);
            compression /= group_overhead;
        }

        // Skip layers reduce effective compression
        if !config.skip_layers.is_empty() {
            let skip_fraction = config.skip_layers.len() as f32 / model.num_layers as f32;
            compression = compression * (1.0 - skip_fraction) + skip_fraction;
        }

        compression.max(1.0)
    }

    /// Convert precision enum to bits.
    fn precision_to_bits(&self, precision: QuantPrecision) -> f32 {
        match precision {
            QuantPrecision::Int2 => 2.0,
            QuantPrecision::Int3 => 3.0,
            QuantPrecision::Int4 => 4.0,
            QuantPrecision::Int8 => 8.0,
            QuantPrecision::FP8 => 8.0,
            QuantPrecision::FP16 => 16.0,
            QuantPrecision::BF16 => 16.0,
            QuantPrecision::Mixed => 6.0, // Average of mixed precision
        }
    }

    /// Get efficiency factor for quantization method.
    fn method_efficiency(&self, method: QuantMethod) -> f32 {
        match method {
            QuantMethod::PTQ => 0.98,  // Post-training has some overhead
            QuantMethod::QAT => 0.99,  // Training-aware is more efficient
            QuantMethod::GPTQ => 0.97, // GPTQ has good compression
            QuantMethod::AWQ => 0.98,  // AWQ is efficient
            QuantMethod::GGML => 0.95, // GGML has metadata overhead
            QuantMethod::SmoothQuant => 0.96,
            QuantMethod::Dynamic => 0.94, // Dynamic has runtime overhead
            QuantMethod::Static => 0.99,
        }
    }

    /// Estimate accuracy degradation from quantization.
    pub fn estimate_degradation(&self, config: &QuantizationConfig) -> f64 {
        let _target_bits = self.precision_to_bits(config.precision);

        // Rough estimate: lower bits = higher degradation
        let base_degradation = match config.precision {
            QuantPrecision::Int2 => 0.15,
            QuantPrecision::Int3 => 0.08,
            QuantPrecision::Int4 => 0.03,
            QuantPrecision::Int8 => 0.01,
            QuantPrecision::FP8 => 0.01,
            QuantPrecision::FP16 => 0.001,
            QuantPrecision::BF16 => 0.002,
            QuantPrecision::Mixed => 0.02,
        };

        // Better methods reduce degradation
        let method_factor = match config.method {
            QuantMethod::GPTQ | QuantMethod::AWQ => 0.7,
            QuantMethod::QAT => 0.5,
            _ => 1.0,
        };

        // Per-channel reduces degradation
        let channel_factor = if config.per_channel { 0.8 } else { 1.0 };

        // Calibration helps
        let calib_factor = 1.0 - (config.calibration_size as f64 / 10000.0).min(0.3);

        base_degradation * method_factor * channel_factor * calib_factor
    }

    /// Simulate calibration data collection.
    pub fn calibrate(&mut self, num_samples: usize) -> Vec<f64> {
        let mut calibration_data = Vec::with_capacity(num_samples);

        for _ in 0..num_samples {
            // Simulate collecting activation statistics
            calibration_data.push(self.random_normal() * 0.1);
        }

        calibration_data
    }

    /// Random number.
    fn random(&mut self) -> f64 {
        self.rng_state = self
            .rng_state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        (self.rng_state >> 11) as f64 / (1u64 << 53) as f64
    }

    /// Random normal.
    fn random_normal(&mut self) -> f64 {
        let u1 = self.random() + 1e-10;
        let u2 = self.random();
        (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
    }
}

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

    #[test]
    fn test_quantizer_creation() {
        let quantizer = Quantizer::new();
        assert_eq!(quantizer.rng_state, 42);
    }

    #[test]
    fn test_quantize_int4() {
        let mut quantizer = Quantizer::new();
        let model = forge_shared::sample_model_info();
        let config = QuantizationConfig {
            precision: QuantPrecision::Int4,
            ..Default::default()
        };

        let compression = quantizer.quantize(&model, &config);
        assert!(compression > 2.0); // FP16 -> INT4 should give ~4x compression
        assert!(compression < 5.0);
    }

    #[test]
    fn test_quantize_int8() {
        let mut quantizer = Quantizer::new();
        let model = forge_shared::sample_model_info();
        let config = QuantizationConfig {
            precision: QuantPrecision::Int8,
            ..Default::default()
        };

        let compression = quantizer.quantize(&model, &config);
        assert!(compression > 1.5); // FP16 -> INT8 should give ~2x compression
        assert!(compression < 2.5);
    }

    #[test]
    fn test_estimate_degradation() {
        let quantizer = Quantizer::new();

        let config_int4 = QuantizationConfig {
            precision: QuantPrecision::Int4,
            ..Default::default()
        };
        let config_int8 = QuantizationConfig {
            precision: QuantPrecision::Int8,
            ..Default::default()
        };

        let deg_int4 = quantizer.estimate_degradation(&config_int4);
        let deg_int8 = quantizer.estimate_degradation(&config_int8);

        // INT4 should have higher degradation than INT8
        assert!(deg_int4 > deg_int8);
    }

    #[test]
    fn test_calibrate() {
        let mut quantizer = Quantizer::new();
        let data = quantizer.calibrate(100);

        assert_eq!(data.len(), 100);
        for val in &data {
            assert!(val.is_finite());
        }
    }

    #[test]
    fn test_group_quantization_overhead() {
        let mut quantizer = Quantizer::new();
        let model = forge_shared::sample_model_info();

        let config_no_group = QuantizationConfig {
            precision: QuantPrecision::Int4,
            group_size: None,
            ..Default::default()
        };
        let config_with_group = QuantizationConfig {
            precision: QuantPrecision::Int4,
            group_size: Some(128),
            ..Default::default()
        };

        let compression_no_group = quantizer.quantize(&model, &config_no_group);
        let compression_with_group = quantizer.quantize(&model, &config_with_group);

        // Group quantization adds overhead, so compression should be slightly less
        assert!(compression_no_group >= compression_with_group);
    }
}