//! Comprehensive tests for quantization functionality
//!
//! Tests cover:
//! - INT8, INT4, and FP8 quantization schemes
//! - Dynamic quantization with calibration
//! - Accuracy validation and loss measurement
//! - Performance optimization verification

use rtx_inference::DataType;
use rtx_inference::quantization::TestTensor;

/// Test quantization configuration
#[tokio::test]
async fn test_quantization_config_creation() {
    use rtx_inference::quantization::{CalibrationMethod, QuantizationConfig, QuantizationScheme};

    let config = QuantizationConfig::new()
        .with_scheme(QuantizationScheme::INT8)
        .with_calibration_method(CalibrationMethod::MinMax)
        .with_calibration_samples(1000)
        .with_accuracy_threshold(0.02)
        .with_dynamic_range_optimization(true);

    assert_eq!(config.scheme, QuantizationScheme::INT8);
    assert_eq!(config.calibration_samples, 1000);
    assert_eq!(config.accuracy_threshold, 0.02);
    assert!(config.dynamic_range_optimization);

    // Test validation
    let invalid_config = QuantizationConfig::new().with_calibration_samples(0); // Should be > 0

    assert!(invalid_config.validate().is_err());
}

/// Test INT8 quantization with symmetric scaling
#[tokio::test]
async fn test_int8_symmetric_quantization() {
    use rtx_inference::quantization::{
        QuantizationConfig, QuantizationScheme, Quantizer, ScalingMethod,
    };

    let config = QuantizationConfig::new()
        .with_scheme(QuantizationScheme::INT8)
        .with_scaling_method(ScalingMethod::Symmetric);

    let quantizer = Quantizer::new(config);

    // Test data with known range
    let input = TestTensor::new(vec![-2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0], vec![8]);

    let quantized = quantizer.quantize(&input).await.unwrap();
    assert_eq!(quantized.dtype, DataType::INT8);

    // Test dequantization
    let dequantized = quantizer.dequantize(&quantized).await.unwrap();
    assert_eq!(dequantized.dtype, DataType::F32);

    // Verify accuracy (should be close to original within quantization error)
    let max_error = input
        .data
        .iter()
        .zip(dequantized.data.iter())
        .map(|(orig, deq)| (orig - deq).abs())
        .fold(0.0f32, |a, b| a.max(b));

    assert!(
        max_error < 0.1,
        "Quantization error too large: {}",
        max_error
    );
}

/// Test INT8 quantization with asymmetric scaling
#[tokio::test]
async fn test_int8_asymmetric_quantization() {
    use rtx_inference::quantization::{
        QuantizationConfig, QuantizationScheme, Quantizer, ScalingMethod,
    };

    let config = QuantizationConfig::new()
        .with_scheme(QuantizationScheme::INT8)
        .with_scaling_method(ScalingMethod::Asymmetric);

    let quantizer = Quantizer::new(config);

    // Test data with asymmetric range
    let input = TestTensor::new(vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], vec![8]);

    let quantized = quantizer.quantize(&input).await.unwrap();
    let dequantized = quantizer.dequantize(&quantized).await.unwrap();

    // Asymmetric quantization should handle this range better
    let max_error = input
        .data
        .iter()
        .zip(dequantized.data.iter())
        .map(|(orig, deq)| (orig - deq).abs())
        .fold(0.0f32, |a, b| a.max(b));

    assert!(
        max_error < 0.05,
        "Asymmetric quantization error: {}",
        max_error
    );
}

/// Test INT4 quantization with weight sharing
#[tokio::test]
async fn test_int4_quantization() {
    use rtx_inference::quantization::{
        GroupSize, QuantizationConfig, QuantizationScheme, Quantizer,
    };

    let config = QuantizationConfig::new()
        .with_scheme(QuantizationScheme::INT4)
        .with_group_size(GroupSize::G128); // Group quantization for better accuracy

    let quantizer = Quantizer::new(config);

    // Create larger tensor for group quantization
    let input = TestTensor::from_random(vec![512, 256], -1.0, 1.0);

    let quantized = quantizer.quantize(&input).await.unwrap();
    assert_eq!(quantized.dtype, DataType::INT4);

    // Verify storage efficiency (INT4 should use ~1/8 the space of F32)
    let compression_ratio =
        quantized.storage_size_bytes() as f32 / input.storage_size_bytes() as f32;
    assert!(
        compression_ratio < 0.15,
        "INT4 compression ratio: {}",
        compression_ratio
    );

    let dequantized = quantizer.dequantize(&quantized).await.unwrap();

    // Measure quantization quality
    let snr = calculate_snr(&input.data, &dequantized.data);
    assert!(snr > 20.0, "Signal-to-noise ratio too low: {} dB", snr);
}

/// Test FP8 E4M3 quantization
#[tokio::test]
async fn test_fp8_e4m3_quantization() {
    use rtx_inference::quantization::{QuantizationConfig, QuantizationScheme, Quantizer};

    let config = QuantizationConfig::new().with_scheme(QuantizationScheme::FP8E4M3);

    let quantizer = Quantizer::new(config);

    // Test with moderate dynamic range (good for E4M3 format)
    let input = TestTensor::from_random(vec![128, 128], -8.0, 8.0);

    let quantized = quantizer.quantize(&input).await.unwrap();
    assert_eq!(quantized.dtype, DataType::FP8E4M3);

    let dequantized = quantizer.dequantize(&quantized).await.unwrap();

    // FP8 should maintain reasonable accuracy for this range
    let mse = calculate_mse(&input.data, &dequantized.data);
    assert!(mse < 0.1, "FP8 E4M3 MSE too high: {}", mse);
}

/// Test FP8 E5M2 quantization
#[tokio::test]
async fn test_fp8_e5m2_quantization() {
    use rtx_inference::quantization::{QuantizationConfig, QuantizationScheme, Quantizer};

    let config = QuantizationConfig::new().with_scheme(QuantizationScheme::FP8E5M2);

    let quantizer = Quantizer::new(config);

    // Test with larger dynamic range (better for E5M2 format)
    let input = TestTensor::from_random(vec![128, 128], -32.0, 32.0);

    let quantized = quantizer.quantize(&input).await.unwrap();
    assert_eq!(quantized.dtype, DataType::FP8E5M2);

    let dequantized = quantizer.dequantize(&quantized).await.unwrap();

    // E5M2 should handle larger ranges better than E4M3
    let relative_error = calculate_relative_error(&input.data, &dequantized.data);
    assert!(
        relative_error < 0.05,
        "FP8 E5M2 relative error: {}",
        relative_error
    );
}

/// Test dynamic quantization with calibration
#[tokio::test]
async fn test_dynamic_quantization_calibration() {
    use rtx_inference::quantization::{CalibrationMethod, DynamicQuantizer, QuantizationConfig};

    let config = QuantizationConfig::new()
        .with_calibration_method(CalibrationMethod::KLDivergence)
        .with_calibration_samples(500);

    let mut quantizer = DynamicQuantizer::new(config);

    // Collect calibration samples
    let mut calibration_data = Vec::new();
    for _ in 0..500 {
        let sample = TestTensor::from_random(vec![64, 64], -10.0, 10.0);
        calibration_data.push(sample);
    }

    // Perform calibration
    quantizer.calibrate(&calibration_data).await.unwrap();

    // Test quantization with calibrated parameters
    let test_input = TestTensor::from_random(vec![64, 64], -8.0, 12.0);
    let quantized = quantizer.quantize(&test_input).await.unwrap();
    let dequantized = quantizer.dequantize(&quantized).await.unwrap();

    // Calibration should improve accuracy
    let calibrated_error = calculate_mse(&test_input.data, &dequantized.data);

    // Compare with uncalibrated quantization
    let uncalibrated_config = QuantizationConfig::new();
    let uncalibrated_quantizer = rtx_inference::quantization::Quantizer::new(uncalibrated_config);
    let uncalibrated_quantized = uncalibrated_quantizer.quantize(&test_input).await.unwrap();
    let uncalibrated_dequantized = uncalibrated_quantizer
        .dequantize(&uncalibrated_quantized)
        .await
        .unwrap();
    let uncalibrated_error = calculate_mse(&test_input.data, &uncalibrated_dequantized.data);

    assert!(
        calibrated_error < uncalibrated_error,
        "Calibrated error ({}) should be less than uncalibrated error ({})",
        calibrated_error,
        uncalibrated_error
    );
}

/// Test accuracy validation with thresholds
#[tokio::test]
async fn test_accuracy_validation() {
    use rtx_inference::quantization::{
        QuantizationConfig, QuantizationValidator, ValidationMetrics,
    };

    let config = QuantizationConfig::new().with_accuracy_threshold(0.01);

    let validator = QuantizationValidator::new(config);

    // Create high-quality quantization (should pass)
    let original = TestTensor::from_random(vec![100, 100], -1.0, 1.0);
    let mut quantized_good = original.clone();
    // Add small noise to simulate good quantization
    for val in quantized_good.data.iter_mut() {
        *val += (fastrand::f32() - 0.5) * 0.005; // Small noise
    }

    let metrics_good = validator
        .validate(&original, &quantized_good)
        .await
        .unwrap();
    assert!(
        metrics_good.passes_threshold(),
        "Good quantization should pass validation"
    );
    assert!(
        metrics_good.mse < 0.01,
        "MSE should be low: {}",
        metrics_good.mse
    );

    // Create poor quantization (should fail)
    let mut quantized_bad = original.clone();
    // Add large noise to simulate poor quantization
    for val in quantized_bad.data.iter_mut() {
        *val += (fastrand::f32() - 0.5) * 0.5; // Large noise
    }

    let metrics_bad = validator.validate(&original, &quantized_bad).await.unwrap();
    assert!(
        !metrics_bad.passes_threshold(),
        "Bad quantization should fail validation"
    );
    assert!(
        metrics_bad.mse > 0.01,
        "MSE should be high: {}",
        metrics_bad.mse
    );
}

/// Test per-channel vs per-tensor quantization
#[tokio::test]
async fn test_per_channel_vs_per_tensor() {
    use rtx_inference::quantization::{QuantizationConfig, QuantizationGranularity, Quantizer};

    let input = TestTensor::new(
        vec![
            // Channel 1: small values
            0.1, 0.2, 0.15, 0.25, // Channel 2: large values
            10.0, 15.0, 12.0, 18.0, // Channel 3: mixed values
            -5.0, 2.0, -1.0, 3.0,
        ],
        vec![3, 4], // 3 channels, 4 elements each
    );

    // Test per-tensor quantization
    let per_tensor_config =
        QuantizationConfig::new().with_granularity(QuantizationGranularity::PerTensor);
    let per_tensor_quantizer = Quantizer::new(per_tensor_config);
    let per_tensor_result = per_tensor_quantizer.quantize(&input).await.unwrap();
    let per_tensor_dequantized = per_tensor_quantizer
        .dequantize(&per_tensor_result)
        .await
        .unwrap();

    // Test per-channel quantization
    let per_channel_config =
        QuantizationConfig::new().with_granularity(QuantizationGranularity::PerChannel);
    let per_channel_quantizer = Quantizer::new(per_channel_config);
    let per_channel_result = per_channel_quantizer.quantize(&input).await.unwrap();
    let per_channel_dequantized = per_channel_quantizer
        .dequantize(&per_channel_result)
        .await
        .unwrap();

    // Per-channel should be more accurate for this case
    let per_tensor_error = calculate_mse(&input.data, &per_tensor_dequantized.data);
    let per_channel_error = calculate_mse(&input.data, &per_channel_dequantized.data);

    assert!(
        per_channel_error < per_tensor_error,
        "Per-channel error ({}) should be less than per-tensor error ({})",
        per_channel_error,
        per_tensor_error
    );
}

/// Test quantization with outlier handling
#[tokio::test]
async fn test_outlier_handling() {
    use rtx_inference::quantization::{OutlierHandling, QuantizationConfig, Quantizer};

    let mut data = vec![1.0; 1000]; // Mostly normal values
    data[500] = 100.0; // Large outlier
    data[501] = -100.0; // Large outlier

    let input = TestTensor::new(data, vec![1000]);

    // Test without outlier handling
    let no_outlier_config = QuantizationConfig::new().with_outlier_handling(OutlierHandling::None);
    let no_outlier_quantizer = Quantizer::new(no_outlier_config);
    let no_outlier_result = no_outlier_quantizer.quantize(&input).await.unwrap();
    let no_outlier_dequantized = no_outlier_quantizer
        .dequantize(&no_outlier_result)
        .await
        .unwrap();

    // Test with outlier clipping
    let clipping_config = QuantizationConfig::new()
        .with_outlier_handling(OutlierHandling::Clipping { percentile: 99.0 });
    let clipping_quantizer = Quantizer::new(clipping_config);
    let clipping_result = clipping_quantizer.quantize(&input).await.unwrap();
    let clipping_dequantized = clipping_quantizer
        .dequantize(&clipping_result)
        .await
        .unwrap();

    // Outlier handling should improve quantization of normal values
    let normal_indices: Vec<usize> = (0..1000).filter(|&i| i != 500 && i != 501).collect();
    let no_outlier_error =
        calculate_mse_subset(&input.data, &no_outlier_dequantized.data, &normal_indices);
    let clipping_error =
        calculate_mse_subset(&input.data, &clipping_dequantized.data, &normal_indices);

    assert!(
        clipping_error < no_outlier_error,
        "Outlier handling should improve normal value accuracy: {} vs {}",
        clipping_error,
        no_outlier_error
    );
}

/// Test quantization performance optimization
#[tokio::test]
async fn test_quantization_performance() {
    use rtx_inference::quantization::{OptimizationLevel, QuantizationConfig, Quantizer};
    use std::time::Instant;

    let large_input = TestTensor::from_random(vec![1024, 1024], -1.0, 1.0);

    // Test basic quantization
    let basic_config = QuantizationConfig::new().with_optimization_level(OptimizationLevel::None);
    let basic_quantizer = Quantizer::new(basic_config);

    let start = Instant::now();
    let _basic_result = basic_quantizer.quantize(&large_input).await.unwrap();
    let basic_time = start.elapsed();

    // Test optimized quantization
    let optimized_config =
        QuantizationConfig::new().with_optimization_level(OptimizationLevel::Aggressive);
    let optimized_quantizer = Quantizer::new(optimized_config);

    let start = Instant::now();
    let _optimized_result = optimized_quantizer.quantize(&large_input).await.unwrap();
    let optimized_time = start.elapsed();

    // Optimized should be faster (though this test might be flaky in CI)
    println!(
        "Basic time: {:?}, Optimized time: {:?}",
        basic_time, optimized_time
    );

    // At least verify both complete successfully
    assert!(basic_time > std::time::Duration::ZERO);
    assert!(optimized_time > std::time::Duration::ZERO);
}

/// Test mixed precision quantization
#[tokio::test]
async fn test_mixed_precision_quantization() {
    use rtx_inference::quantization::{LayerConfig, MixedPrecisionQuantizer, QuantizationScheme};

    let mut quantizer = MixedPrecisionQuantizer::new();

    // Configure different layers with different precisions
    quantizer.add_layer_config("embedding", LayerConfig::new(QuantizationScheme::INT8));
    quantizer.add_layer_config("attention", LayerConfig::new(QuantizationScheme::FP8E4M3));
    quantizer.add_layer_config("ffn", LayerConfig::new(QuantizationScheme::INT4));
    quantizer.add_layer_config("output", LayerConfig::new(QuantizationScheme::FP8E5M2));

    // Test each layer type
    let test_cases = vec![
        (
            "embedding",
            TestTensor::from_random(vec![512, 768], -0.1, 0.1),
        ),
        (
            "attention",
            TestTensor::from_random(vec![64, 64], -2.0, 2.0),
        ),
        ("ffn", TestTensor::from_random(vec![768, 3072], -1.0, 1.0)),
        (
            "output",
            TestTensor::from_random(vec![3072, 50257], -0.5, 0.5),
        ),
    ];

    for (layer_name, input) in test_cases {
        let quantized = quantizer.quantize_layer(layer_name, &input).await.unwrap();
        let dequantized = quantizer
            .dequantize_layer(layer_name, &quantized)
            .await
            .unwrap();

        // Verify that appropriate precision was used
        let expected_dtype = match layer_name {
            "embedding" => DataType::INT8,
            "attention" => DataType::FP8E4M3,
            "ffn" => DataType::INT4,
            "output" => DataType::FP8E5M2,
            _ => panic!("Unexpected layer name"),
        };
        assert_eq!(quantized.dtype, expected_dtype);

        // Verify reasonable accuracy
        let error = calculate_mse(&input.data, &dequantized.data);
        assert!(
            error < 0.5,
            "Layer {} error too high: {}",
            layer_name,
            error
        );
    }
}

/// Test error handling in quantization
#[tokio::test]
async fn test_quantization_error_handling() {
    use rtx_inference::quantization::{QuantizationConfig, QuantizationScheme, Quantizer};

    let quantizer = Quantizer::new(QuantizationConfig::new());

    // Test empty tensor
    let empty_tensor = TestTensor::new(vec![], vec![0]);
    let result = quantizer.quantize(&empty_tensor).await;
    assert!(result.is_err(), "Should fail on empty tensor");

    // Test invalid shape
    let invalid_tensor = TestTensor::new(vec![1.0, 2.0], vec![3]); // Data/shape mismatch
    let result = quantizer.quantize(&invalid_tensor).await;
    assert!(result.is_err(), "Should fail on shape mismatch");

    // Test extreme values
    let extreme_tensor = TestTensor::new(vec![f32::INFINITY, f32::NEG_INFINITY, f32::NAN], vec![3]);
    let result = quantizer.quantize(&extreme_tensor).await;
    assert!(result.is_err(), "Should fail on extreme values");
}

/// Test quantization with batch processing
#[tokio::test]
async fn test_batch_quantization() {
    use rtx_inference::quantization::{BatchQuantizer, QuantizationConfig};

    let config = QuantizationConfig::new();
    let quantizer = BatchQuantizer::new(config);

    // Create batch of tensors
    let batch = vec![
        TestTensor::from_random(vec![128, 256], -1.0, 1.0),
        TestTensor::from_random(vec![128, 256], -2.0, 2.0),
        TestTensor::from_random(vec![128, 256], -0.5, 0.5),
        TestTensor::from_random(vec![128, 256], -3.0, 3.0),
    ];

    let quantized_batch = quantizer.quantize_batch(&batch).await.unwrap();
    assert_eq!(quantized_batch.len(), batch.len());

    let dequantized_batch = quantizer.dequantize_batch(&quantized_batch).await.unwrap();
    assert_eq!(dequantized_batch.len(), batch.len());

    // Verify accuracy for each tensor in batch
    for (i, (original, dequantized)) in batch.iter().zip(dequantized_batch.iter()).enumerate() {
        let error = calculate_mse(&original.data, &dequantized.data);
        assert!(error < 0.1, "Batch item {} error too high: {}", i, error);
    }
}

// Helper functions

fn calculate_snr(original: &[f32], quantized: &[f32]) -> f32 {
    let signal_power: f32 = original.iter().map(|x| x * x).sum();
    let noise_power: f32 = original
        .iter()
        .zip(quantized.iter())
        .map(|(o, q)| (o - q) * (o - q))
        .sum();

    if noise_power == 0.0 {
        f32::INFINITY
    } else {
        10.0 * (signal_power / noise_power).log10()
    }
}

fn calculate_mse(original: &[f32], quantized: &[f32]) -> f32 {
    assert_eq!(original.len(), quantized.len());
    let sum_squared_error: f32 = original
        .iter()
        .zip(quantized.iter())
        .map(|(o, q)| (o - q) * (o - q))
        .sum();
    sum_squared_error / original.len() as f32
}

fn calculate_relative_error(original: &[f32], quantized: &[f32]) -> f32 {
    assert_eq!(original.len(), quantized.len());
    let sum_relative_error: f32 = original
        .iter()
        .zip(quantized.iter())
        .map(|(o, q)| {
            if o.abs() < 1e-8 {
                if q.abs() < 1e-8 { 0.0 } else { 1.0 }
            } else {
                ((o - q) / o).abs()
            }
        })
        .sum();
    sum_relative_error / original.len() as f32
}

fn calculate_mse_subset(original: &[f32], quantized: &[f32], indices: &[usize]) -> f32 {
    let sum_squared_error: f32 = indices
        .iter()
        .map(|&i| (original[i] - quantized[i]) * (original[i] - quantized[i]))
        .sum();
    sum_squared_error / indices.len() as f32
}