//! Comprehensive tests for post-training quantization
//!
//! Tests cover:
//! - INT8/INT4/INT2/INT1 quantization schemes
//! - Calibration dataset handling
#![cfg(feature = "disabled_tests")]
//! - Accuracy preservation metrics
//! - Performance benchmarking
//! - Edge cases and error conditions

use rtx_compress::{
    CompressionError, Result,
    quantization::{
        CalibrationConfig, PostTrainingQuantizer, QuantizationConfig, QuantizationError,
        QuantizationScheme, QuantizationStatistics,
    },
};
use rtx_tensor::{Device, Tensor};
use std::collections::HashMap;

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

    fn create_test_tensor(shape: &[usize]) -> Result<Tensor> {
        let device = Device::try_default()?;
        // Create a tensor with known distribution for testing
        let mut data = Vec::new();
        let total_elements: usize = shape.iter().product();

        for i in 0..total_elements {
            // Create a mix of small and large values to test quantization
            let val = match i % 4 {
                0 => (i as f32) / 1000.0,  // Small positive values
                1 => -(i as f32) / 1000.0, // Small negative values
                2 => (i as f32) * 0.1,     // Larger positive values
                _ => -(i as f32) * 0.1,    // Larger negative values
            };
            data.push(val);
        }

        Tensor::from_slice(&data, shape, &device)
    }

    fn create_calibration_dataset(num_samples: usize, shape: &[usize]) -> Result<Vec<Tensor>> {
        let mut dataset = Vec::new();
        for _ in 0..num_samples {
            dataset.push(create_test_tensor(shape)?);
        }
        Ok(dataset)
    }

    #[test]
    fn test_int8_post_training_quantization() -> Result<()> {
        let config = QuantizationConfig::new(
            QuantizationScheme::INT8,
            CalibrationConfig::new(100, 0.99), // 100 samples, 99% coverage
        );

        let quantizer = PostTrainingQuantizer::new(config)?;

        // Create test tensor and calibration data
        let tensor = create_test_tensor(&[64, 128])?;
        let calibration_data = create_calibration_dataset(100, &[64, 128])?;

        // Calibrate quantizer
        let mut quantizer = quantizer;
        quantizer.calibrate(&calibration_data)?;

        // Quantize tensor
        let quantized = quantizer.quantize(&tensor)?;
        let dequantized = quantizer.dequantize(&quantized)?;

        // Verify shapes are preserved
        assert_eq!(tensor.shape(), dequantized.shape());

        // Verify quantization statistics
        let stats = quantizer.get_statistics();
        assert!(stats.scale > 0.0);
        assert!(stats.zero_point >= -128 && stats.zero_point <= 127);
        assert_eq!(stats.bit_width, 8);

        // Verify accuracy is within acceptable bounds (should be high for INT8)
        let mse = calculate_mse(&tensor, &dequantized)?;
        assert!(mse < 0.01, "MSE too high: {}", mse);

        Ok(())
    }

    #[test]
    fn test_int4_post_training_quantization() -> Result<()> {
        let config = QuantizationConfig::new(
            QuantizationScheme::INT4,
            CalibrationConfig::new(100, 0.95), // Slightly lower coverage for more aggressive quantization
        );

        let quantizer = PostTrainingQuantizer::new(config)?;
        let tensor = create_test_tensor(&[32, 64])?;
        let calibration_data = create_calibration_dataset(100, &[32, 64])?;

        let mut quantizer = quantizer;
        quantizer.calibrate(&calibration_data)?;

        let quantized = quantizer.quantize(&tensor)?;
        let dequantized = quantizer.dequantize(&quantized)?;

        assert_eq!(tensor.shape(), dequantized.shape());

        let stats = quantizer.get_statistics();
        assert_eq!(stats.bit_width, 4);
        assert!(stats.zero_point >= -8 && stats.zero_point <= 7);

        // INT4 will have higher error than INT8
        let mse = calculate_mse(&tensor, &dequantized)?;
        assert!(mse < 0.1, "MSE too high: {}", mse);

        Ok(())
    }

    #[test]
    fn test_int2_post_training_quantization() -> Result<()> {
        let config =
            QuantizationConfig::new(QuantizationScheme::INT2, CalibrationConfig::new(50, 0.9));

        let quantizer = PostTrainingQuantizer::new(config)?;
        let tensor = create_test_tensor(&[16, 32])?;
        let calibration_data = create_calibration_dataset(50, &[16, 32])?;

        let mut quantizer = quantizer;
        quantizer.calibrate(&calibration_data)?;

        let quantized = quantizer.quantize(&tensor)?;
        let dequantized = quantizer.dequantize(&quantized)?;

        assert_eq!(tensor.shape(), dequantized.shape());

        let stats = quantizer.get_statistics();
        assert_eq!(stats.bit_width, 2);
        assert!(stats.zero_point >= -2 && stats.zero_point <= 1);

        // INT2 will have much higher error
        let mse = calculate_mse(&tensor, &dequantized)?;
        assert!(mse < 1.0, "MSE too high: {}", mse);

        Ok(())
    }

    #[test]
    fn test_int1_post_training_quantization() -> Result<()> {
        let config = QuantizationConfig::new(
            QuantizationScheme::INT1, // Binary quantization
            CalibrationConfig::new(50, 0.8),
        );

        let quantizer = PostTrainingQuantizer::new(config)?;
        let tensor = create_test_tensor(&[8, 16])?;
        let calibration_data = create_calibration_dataset(50, &[8, 16])?;

        let mut quantizer = quantizer;
        quantizer.calibrate(&calibration_data)?;

        let quantized = quantizer.quantize(&tensor)?;
        let dequantized = quantizer.dequantize(&quantized)?;

        assert_eq!(tensor.shape(), dequantized.shape());

        let stats = quantizer.get_statistics();
        assert_eq!(stats.bit_width, 1);
        // For binary quantization, values should be either -1 or 1
        assert!(stats.zero_point == -1 || stats.zero_point == 0 || stats.zero_point == 1);

        // Binary quantization will have the highest error
        let mse = calculate_mse(&tensor, &dequantized)?;
        assert!(mse < 10.0, "MSE too high: {}", mse);

        Ok(())
    }

    #[test]
    fn test_calibration_with_insufficient_data() -> Result<()> {
        let config = QuantizationConfig::new(
            QuantizationScheme::INT8,
            CalibrationConfig::new(100, 0.99), // Requires 100 samples
        );

        let quantizer = PostTrainingQuantizer::new(config)?;
        let calibration_data = create_calibration_dataset(50, &[32, 64])?; // Only 50 samples

        let mut quantizer = quantizer;
        let result = quantizer.calibrate(&calibration_data);

        // Should handle insufficient data gracefully
        match result {
            Ok(_) => {
                // Implementation should work with available data
                let stats = quantizer.get_statistics();
                assert!(stats.calibration_samples < 100);
            }
            Err(CompressionError::Quantization(QuantizationError::InsufficientCalibrationData)) => {
                // Or return appropriate error
            }
            Err(e) => panic!("Unexpected error: {:?}", e),
        }

        Ok(())
    }

    #[test]
    fn test_quantization_without_calibration() {
        let config =
            QuantizationConfig::new(QuantizationScheme::INT8, CalibrationConfig::new(100, 0.99));

        let quantizer = PostTrainingQuantizer::new(config).unwrap();
        let tensor = create_test_tensor(&[32, 64]).unwrap();

        // Should fail without calibration
        let result = quantizer.quantize(&tensor);
        assert!(result.is_err());

        match result {
            Err(CompressionError::Quantization(QuantizationError::NotCalibrated)) => {
                // Expected error
            }
            Err(e) => panic!("Unexpected error: {:?}", e),
            Ok(_) => panic!("Should have failed without calibration"),
        }
    }

    #[test]
    fn test_percentile_based_calibration() -> Result<()> {
        let config = QuantizationConfig::new(
            QuantizationScheme::INT8,
            CalibrationConfig::new_with_percentile(100, 0.01, 0.99), // Use 1st to 99th percentile
        );

        let quantizer = PostTrainingQuantizer::new(config)?;
        let calibration_data = create_calibration_dataset(100, &[32, 64])?;

        let mut quantizer = quantizer;
        quantizer.calibrate(&calibration_data)?;

        let stats = quantizer.get_statistics();
        assert!(stats.min_value < stats.max_value);
        assert!(stats.calibration_method == "percentile");

        Ok(())
    }

    #[test]
    fn test_entropy_based_calibration() -> Result<()> {
        let config = QuantizationConfig::new(
            QuantizationScheme::INT8,
            CalibrationConfig::new_entropy_based(100, 128), // 128 histogram bins
        );

        let quantizer = PostTrainingQuantizer::new(config)?;
        let calibration_data = create_calibration_dataset(100, &[32, 64])?;

        let mut quantizer = quantizer;
        quantizer.calibrate(&calibration_data)?;

        let stats = quantizer.get_statistics();
        assert_eq!(stats.calibration_method, "entropy");
        assert!(stats.entropy_score > 0.0);

        Ok(())
    }

    #[test]
    fn test_symmetric_vs_asymmetric_quantization() -> Result<()> {
        // Test symmetric quantization
        let symmetric_config = QuantizationConfig::new_symmetric(
            QuantizationScheme::INT8,
            CalibrationConfig::new(100, 0.99),
        );

        let mut symmetric_quantizer = PostTrainingQuantizer::new(symmetric_config)?;
        let calibration_data = create_calibration_dataset(100, &[32, 64])?;
        symmetric_quantizer.calibrate(&calibration_data)?;

        let symmetric_stats = symmetric_quantizer.get_statistics();
        assert_eq!(symmetric_stats.zero_point, 0); // Symmetric should have zero_point = 0

        // Test asymmetric quantization
        let asymmetric_config = QuantizationConfig::new_asymmetric(
            QuantizationScheme::INT8,
            CalibrationConfig::new(100, 0.99),
        );

        let mut asymmetric_quantizer = PostTrainingQuantizer::new(asymmetric_config)?;
        asymmetric_quantizer.calibrate(&calibration_data)?;

        let asymmetric_stats = asymmetric_quantizer.get_statistics();
        // Asymmetric may have non-zero zero_point
        assert!(asymmetric_stats.zero_point >= -128 && asymmetric_stats.zero_point <= 127);

        Ok(())
    }

    #[test]
    fn test_per_channel_quantization() -> Result<()> {
        let config = QuantizationConfig::new_per_channel(
            QuantizationScheme::INT8,
            CalibrationConfig::new(100, 0.99),
            0, // Quantize along channel dimension 0
        );

        let quantizer = PostTrainingQuantizer::new(config)?;
        let tensor = create_test_tensor(&[64, 128])?;
        let calibration_data = create_calibration_dataset(100, &[64, 128])?;

        let mut quantizer = quantizer;
        quantizer.calibrate(&calibration_data)?;

        let quantized = quantizer.quantize(&tensor)?;
        let dequantized = quantizer.dequantize(&quantized)?;

        assert_eq!(tensor.shape(), dequantized.shape());

        let stats = quantizer.get_statistics();
        assert_eq!(stats.quantization_type, "per_channel");
        assert_eq!(stats.num_channels, 64); // Should match first dimension

        // Per-channel quantization should have better accuracy than per-tensor
        let mse = calculate_mse(&tensor, &dequantized)?;
        assert!(
            mse < 0.005,
            "MSE too high for per-channel quantization: {}",
            mse
        );

        Ok(())
    }

    #[test]
    fn test_quantization_compression_ratio() -> Result<()> {
        let config =
            QuantizationConfig::new(QuantizationScheme::INT8, CalibrationConfig::new(100, 0.99));

        let quantizer = PostTrainingQuantizer::new(config)?;
        let tensor = create_test_tensor(&[64, 128])?;
        let calibration_data = create_calibration_dataset(100, &[64, 128])?;

        let mut quantizer = quantizer;
        quantizer.calibrate(&calibration_data)?;

        let quantized = quantizer.quantize(&tensor)?;

        // Calculate compression ratio
        let original_size = tensor.numel() * 4; // f32 = 4 bytes
        let compressed_size = quantized.storage_size();
        let compression_ratio = original_size as f32 / compressed_size as f32;

        // INT8 should provide ~4x compression ratio
        assert!(
            compression_ratio > 3.5 && compression_ratio < 4.5,
            "Unexpected compression ratio: {}",
            compression_ratio
        );

        Ok(())
    }

    #[test]
    fn test_batch_quantization() -> Result<()> {
        let config =
            QuantizationConfig::new(QuantizationScheme::INT8, CalibrationConfig::new(100, 0.99));

        let quantizer = PostTrainingQuantizer::new(config)?;
        let calibration_data = create_calibration_dataset(100, &[32, 64])?;

        let mut quantizer = quantizer;
        quantizer.calibrate(&calibration_data)?;

        // Create batch of tensors
        let batch_tensors = vec![
            create_test_tensor(&[32, 64])?,
            create_test_tensor(&[32, 64])?,
            create_test_tensor(&[32, 64])?,
        ];

        let quantized_batch = quantizer.quantize_batch(&batch_tensors)?;
        let dequantized_batch = quantizer.dequantize_batch(&quantized_batch)?;

        assert_eq!(batch_tensors.len(), dequantized_batch.len());

        for (original, dequantized) in batch_tensors.iter().zip(dequantized_batch.iter()) {
            assert_eq!(original.shape(), dequantized.shape());
            let mse = calculate_mse(original, dequantized)?;
            assert!(mse < 0.01, "MSE too high: {}", mse);
        }

        Ok(())
    }

    #[test]
    fn test_quantization_serialization() -> Result<()> {
        let config =
            QuantizationConfig::new(QuantizationScheme::INT8, CalibrationConfig::new(100, 0.99));

        let quantizer = PostTrainingQuantizer::new(config)?;
        let calibration_data = create_calibration_dataset(100, &[32, 64])?;

        let mut quantizer = quantizer;
        quantizer.calibrate(&calibration_data)?;

        // Serialize quantizer
        let serialized = quantizer.serialize()?;

        // Deserialize quantizer
        let deserialized_quantizer = PostTrainingQuantizer::deserialize(&serialized)?;

        // Verify they produce the same results
        let tensor = create_test_tensor(&[32, 64])?;
        let original_quantized = quantizer.quantize(&tensor)?;
        let deserialized_quantized = deserialized_quantizer.quantize(&tensor)?;

        assert_eq!(original_quantized.shape(), deserialized_quantized.shape());
        // Note: For testing, we'd need to implement proper comparison of quantized tensors

        Ok(())
    }

    // Helper function to calculate MSE between two tensors
    fn calculate_mse(tensor1: &Tensor, tensor2: &Tensor) -> Result<f32> {
        let diff = tensor1.sub(tensor2)?;
        let squared_diff = diff.mul(&diff)?;
        let mse = squared_diff.mean(None, false)?.to_scalar::<f32>()?;
        Ok(mse)
    }

    #[test]
    fn test_mixed_precision_quantization() -> Result<()> {
        // Test quantizing different layers with different precision
        let mut layer_configs = HashMap::new();
        layer_configs.insert("conv1".to_string(), QuantizationScheme::INT8);
        layer_configs.insert("conv2".to_string(), QuantizationScheme::INT4);
        layer_configs.insert("fc".to_string(), QuantizationScheme::INT8); // Keep final layer higher precision

        let config = QuantizationConfig::new_mixed_precision(
            layer_configs,
            CalibrationConfig::new(100, 0.99),
        );

        let quantizer = PostTrainingQuantizer::new(config)?;
        let calibration_data = create_calibration_dataset(100, &[32, 64])?;

        let mut quantizer = quantizer;
        quantizer.calibrate(&calibration_data)?;

        // Create tensors for different layers
        let conv1_tensor = create_test_tensor(&[32, 64])?;
        let conv2_tensor = create_test_tensor(&[32, 64])?;
        let fc_tensor = create_test_tensor(&[32, 64])?;

        let conv1_quantized = quantizer.quantize_layer(&conv1_tensor, "conv1")?;
        let conv2_quantized = quantizer.quantize_layer(&conv2_tensor, "conv2")?;
        let fc_quantized = quantizer.quantize_layer(&fc_tensor, "fc")?;

        // Verify different compression ratios
        let conv1_ratio = calculate_compression_ratio(&conv1_tensor, &conv1_quantized);
        let conv2_ratio = calculate_compression_ratio(&conv2_tensor, &conv2_quantized);
        let fc_ratio = calculate_compression_ratio(&fc_tensor, &fc_quantized);

        assert!(conv2_ratio > conv1_ratio); // INT4 should compress more than INT8
        assert!((conv1_ratio - fc_ratio).abs() < 0.1); // Both INT8 should be similar

        Ok(())
    }

    fn calculate_compression_ratio(original: &Tensor, compressed: &impl StorageSized) -> f32 {
        let original_size = original.numel() * 4; // f32 = 4 bytes
        let compressed_size = compressed.storage_size();
        original_size as f32 / compressed_size as f32
    }
}

// Trait for objects that have a storage size
trait StorageSized {
    fn storage_size(&self) -> usize;
}

// This would need to be implemented for the actual quantized tensor type
impl StorageSized for Tensor {
    fn storage_size(&self) -> usize {
        self.numel() * 4 // Placeholder - would depend on actual storage
    }
}