//! Comprehensive tests for quantization functionality

use rtx_compress::{
    CompressionError, Result,
    quantization::{
        CalibrationConfig, CalibrationMethod, PQConfig, PostTrainingQuantizer, ProductQuantizer,
        QuantizationConfig, QuantizationScheme,
    },
};
use rtx_tensor::{Device, Tensor};

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

    fn create_test_device() -> Device {
        Device::cpu()
    }

    fn create_test_tensor(size: usize) -> Tensor {
        let data: Vec<f32> = (0..size).map(|i| i as f32 / 10.0).collect();
        Tensor::from_data(data, &[size / 4, 4], &create_test_device()).unwrap()
    }

    fn create_embeddings_tensor() -> Tensor {
        // Create some embeddings for testing [batch_size=8, embedding_dim=128]
        let data: Vec<f32> = (0..1024).map(|i| (i as f32).sin()).collect();
        Tensor::from_data(data, &[8, 128], &create_test_device()).unwrap()
    }

    #[test]
    fn test_pq_config_creation() {
        let config = PQConfig {
            num_subquantizers: 8,
            codebook_size: 256,
            max_iterations: 100,
            tolerance: 1e-4,
            use_opq: false,
            opq_iterations: 20,
            use_residual: false,
            residual_stages: 2,
        };

        assert_eq!(config.num_subquantizers, 8);
        assert_eq!(config.codebook_size, 256);
        assert_eq!(config.max_iterations, 100);
        assert!(!config.use_opq);
    }

    #[test]
    fn test_product_quantizer_creation() {
        let config = PQConfig {
            num_subquantizers: 4,
            codebook_size: 256,
            max_iterations: 50,
            tolerance: 1e-4,
            use_opq: false,
            opq_iterations: 20,
            use_residual: false,
            residual_stages: 2,
        };

        let quantizer = ProductQuantizer::new(config);
        assert!(quantizer.is_ok());
    }

    #[test]
    #[ignore = "Pre-existing quantizer fit issue"]
    fn test_product_quantizer_train() {
        let config = PQConfig {
            num_subquantizers: 4,
            codebook_size: 16, // Small for testing
            max_iterations: 10,
            tolerance: 1e-3,
            use_opq: false,
            opq_iterations: 20,
            use_residual: false,
            residual_stages: 2,
        };

        let mut quantizer = ProductQuantizer::new(config).unwrap();
        let embeddings = create_embeddings_tensor();

        let result = quantizer.fit(&embeddings);
        assert!(result.is_ok());
    }

    #[test]
    #[ignore = "Pre-existing codebook size issue"]
    fn test_product_quantizer_encode_decode() {
        let config = PQConfig {
            num_subquantizers: 4,
            codebook_size: 16,
            max_iterations: 10,
            tolerance: 1e-3,
            use_opq: false,
            opq_iterations: 20,
            use_residual: false,
            residual_stages: 2,
        };

        let mut quantizer = ProductQuantizer::new(config).unwrap();
        let embeddings = create_embeddings_tensor();

        // Train the quantizer
        quantizer.fit(&embeddings).unwrap();

        // Encode
        let encoded = quantizer.encode(&embeddings);
        assert!(encoded.is_ok());

        // Decode
        let decoded = quantizer.decode(&encoded.unwrap());
        assert!(decoded.is_ok());

        // Check shape is preserved
        assert_eq!(decoded.unwrap().shape(), embeddings.shape());
    }

    #[test]
    fn test_post_training_quantizer() {
        let calibration = CalibrationConfig::new(10, 1.0);
        let config = QuantizationConfig::new_symmetric(QuantizationScheme::INT8, calibration);

        let quantizer = PostTrainingQuantizer::new(config);
        assert!(quantizer.is_ok());
    }

    #[test]
    fn test_int8_quantization() {
        let calibration = CalibrationConfig::new(1, 1.0);
        let config = QuantizationConfig::new_symmetric(QuantizationScheme::INT8, calibration);

        let mut quantizer = PostTrainingQuantizer::new(config).unwrap();
        let tensor = create_test_tensor(64);

        // Calibrate
        quantizer.calibrate(&[tensor.clone()]).unwrap();

        // Quantize
        let quantized = quantizer.quantize(&tensor);
        assert!(quantized.is_ok());

        // Dequantize
        let dequantized = quantizer.dequantize(&quantized.unwrap());
        assert!(dequantized.is_ok());
    }

    #[test]
    fn test_int4_quantization() {
        let calibration = CalibrationConfig::new(1, 1.0);
        let config = QuantizationConfig::new_asymmetric(QuantizationScheme::INT4, calibration);

        let mut quantizer = PostTrainingQuantizer::new(config).unwrap();
        let tensor = create_test_tensor(64);

        quantizer.calibrate(&[tensor.clone()]).unwrap();
        let quantized = quantizer.quantize(&tensor).unwrap();
        let dequantized = quantizer.dequantize(&quantized).unwrap();

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

    #[test]
    fn test_percentile_calibration() {
        let calibration = CalibrationConfig::new_with_percentile(1, 0.001, 0.999);
        let config = QuantizationConfig::new_symmetric(QuantizationScheme::INT8, calibration);

        let mut quantizer = PostTrainingQuantizer::new(config).unwrap();
        let tensor = create_test_tensor(100);

        let result = quantizer.calibrate(&[tensor]);
        assert!(result.is_ok());
    }

    #[test]
    fn test_per_channel_quantization() {
        let calibration = CalibrationConfig::new(1, 1.0);
        let config = QuantizationConfig::new_per_channel(
            QuantizationScheme::INT8,
            calibration,
            0, // channel dimension
        );

        let mut quantizer = PostTrainingQuantizer::new(config).unwrap();
        let tensor = create_test_tensor(64);

        quantizer.calibrate(&[tensor.clone()]).unwrap();
        let quantized = quantizer.quantize(&tensor).unwrap();
        let dequantized = quantizer.dequantize(&quantized).unwrap();

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

    #[test]
    #[ignore = "Pre-existing codebook size issue"]
    fn test_compression_ratio() {
        let config = PQConfig {
            num_subquantizers: 8,
            codebook_size: 256, // 8 bits per subquantizer
            max_iterations: 10,
            tolerance: 1e-3,
            use_opq: false,
            opq_iterations: 20,
            use_residual: false,
            residual_stages: 2,
        };

        let mut quantizer = ProductQuantizer::new(config).unwrap();
        let embeddings = create_embeddings_tensor();

        quantizer.fit(&embeddings).unwrap();
        let encoded = quantizer.encode(&embeddings).unwrap();

        // Calculate compression ratio
        let original_size = embeddings.numel() * 4; // f32 = 4 bytes
        let compressed_size = encoded.numel(); // indices are typically u8 or u16

        let ratio = original_size as f64 / compressed_size as f64;
        assert!(ratio > 1.0); // Should achieve compression
    }

    #[test]
    #[ignore = "Pre-existing codebook size issue"]
    fn test_residual_quantization() {
        let config = PQConfig {
            num_subquantizers: 4,
            codebook_size: 16,
            max_iterations: 10,
            tolerance: 1e-3,
            use_opq: false,
            opq_iterations: 20,
            use_residual: true,
            residual_stages: 2,
        };

        let mut quantizer = ProductQuantizer::new(config).unwrap();
        let embeddings = create_embeddings_tensor();

        quantizer.fit(&embeddings).unwrap();
        let encoded = quantizer.encode(&embeddings).unwrap();
        let decoded = quantizer.decode(&encoded).unwrap();

        // With residual quantization, reconstruction should be better
        assert_eq!(decoded.shape(), embeddings.shape());
    }

    #[test]
    fn test_quantization_statistics() {
        let calibration = CalibrationConfig::new(1, 1.0);
        let config = QuantizationConfig::new_symmetric(QuantizationScheme::INT8, calibration);

        let mut quantizer = PostTrainingQuantizer::new(config).unwrap();
        let tensor = create_test_tensor(64);

        quantizer.calibrate(&[tensor.clone()]).unwrap();
        let quantized = quantizer.quantize(&tensor).unwrap();

        let stats = quantizer.get_statistics();
        assert!(stats.scale > 0.0);
        assert!(stats.bit_width > 0);
    }
}