//! Per-Channel Quantization Tests
//!
//! Validates that per-channel quantization reduces quantization error from 2.5% to 1.5%
//! on attention weights and linear layer weights.

use candle_core::{DType, Device, Tensor};
use ml::memory_optimization::quantization::{QuantizationConfig, QuantizationType, Quantizer};
use ml::MLError;

/// Test 1: Per-channel quantization reduces error on attention weights
#[test]
fn test_per_channel_attention_weight_accuracy() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Create attention weight matrix [256, 256] (out_channels, in_channels)
    let weight_data: Vec<f32> = (0..256 * 256)
        .map(|i| ((i as f32) * 0.01).sin() * 0.5)
        .collect();
    let weight = Tensor::from_slice(&weight_data, (256, 256), &device)?;

    // Test per-tensor quantization (per_channel = false)
    let config_per_tensor = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer_per_tensor = Quantizer::new(config_per_tensor, device.clone());
    let quantized_per_tensor = quantizer_per_tensor.quantize_tensor(&weight, "q_weight")?;
    let dequantized_per_tensor = quantizer_per_tensor.dequantize_tensor(&quantized_per_tensor)?;

    // Calculate per-tensor error
    let error_per_tensor = calculate_relative_error(&weight, &dequantized_per_tensor)?;

    // Test per-channel quantization (per_channel = true)
    let config_per_channel = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer_per_channel = Quantizer::new(config_per_channel, device.clone());
    let quantized_per_channel = quantizer_per_channel.quantize_tensor(&weight, "q_weight")?;

    // Verify per-channel params were stored
    assert!(
        quantizer_per_channel.has_per_channel_params("q_weight"),
        "Per-channel params should be stored"
    );

    // Dequantize using per-channel method
    let dequantized_per_channel =
        quantizer_per_channel.dequantize_tensor_per_channel(&quantized_per_channel, "q_weight")?;

    // Calculate per-channel error
    let error_per_channel = calculate_relative_error(&weight, &dequantized_per_channel)?;

    println!(
        "Per-tensor error: {:.4}%, Per-channel error: {:.4}%",
        error_per_tensor * 100.0,
        error_per_channel * 100.0
    );

    // Validate error reduction: per-channel should be lower than per-tensor
    assert!(
        error_per_channel < error_per_tensor,
        "Per-channel error {:.4}% should be lower than per-tensor error {:.4}%",
        error_per_channel * 100.0,
        error_per_tensor * 100.0
    );

    // Target validation: per-channel error should be < 1.5%
    assert!(
        error_per_channel < 0.015,
        "Per-channel error {:.4}% should be < 1.5%",
        error_per_channel * 100.0
    );

    Ok(())
}

/// Test 2: Per-channel quantization on linear layer weights
#[test]
fn test_per_channel_linear_layer_accuracy() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Create linear layer weight [128, 256] (out_features, in_features)
    let weight_data: Vec<f32> = (0..128 * 256)
        .map(|i| ((i as f32) * 0.02).cos() * 0.3)
        .collect();
    let weight = Tensor::from_slice(&weight_data, (128, 256), &device)?;

    // Per-channel quantization
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());

    // Quantize
    let quantized = quantizer.quantize_tensor(&weight, "linear_weight")?;

    // Verify quantized dtype is U8
    assert_eq!(
        quantized.data.dtype(),
        DType::U8,
        "Quantized data should be U8"
    );

    // Verify shape is preserved
    assert_eq!(
        quantized.data.dims(),
        &[128, 256],
        "Shape should be preserved"
    );

    // Dequantize
    let dequantized = quantizer.dequantize_tensor_per_channel(&quantized, "linear_weight")?;

    // Calculate error
    let error = calculate_relative_error(&weight, &dequantized)?;

    println!("Linear layer per-channel error: {:.4}%", error * 100.0);

    // Validate error < 1.5%
    assert!(
        error < 0.015,
        "Per-channel error {:.4}% should be < 1.5%",
        error * 100.0
    );

    Ok(())
}

/// Test 3: Per-channel parameters storage and retrieval
#[test]
fn test_per_channel_params_storage() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Create weight [64, 128]
    let weight_data: Vec<f32> = (0..64 * 128).map(|i| (i as f32) * 0.01).collect();
    let weight = Tensor::from_slice(&weight_data, (64, 128), &device)?;

    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());

    // Quantize
    let _quantized = quantizer.quantize_tensor(&weight, "test_weight")?;

    // Check params were stored
    assert!(
        quantizer.has_per_channel_params("test_weight"),
        "Per-channel params should be stored"
    );

    // Retrieve params
    let params = quantizer
        .get_per_channel_params("test_weight")
        .expect("Params should exist");

    // Verify params shape
    assert_eq!(
        params.scales.len(),
        64,
        "Should have 64 scales (one per output channel)"
    );
    assert_eq!(params.zero_points.len(), 64, "Should have 64 zero points");
    assert_eq!(params.min_vals.len(), 64, "Should have 64 min values");
    assert_eq!(params.max_vals.len(), 64, "Should have 64 max values");

    // Verify scales are positive
    for (idx, scale) in params.scales.iter().enumerate() {
        assert!(
            *scale > 0.0,
            "Scale {} should be positive, got {}",
            idx,
            scale
        );
    }

    // Verify zero points are within valid range for symmetric quantization
    for (idx, zp) in params.zero_points.iter().enumerate() {
        assert_eq!(
            *zp, 127,
            "Zero point {} should be 127 (symmetric), got {}",
            idx, zp
        );
    }

    Ok(())
}

/// Test 4: Per-channel quantization comparison with per-tensor
#[test]
fn test_per_channel_vs_per_tensor_comparison() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Create weight with varying scales across channels
    // Each channel has different magnitude to amplify per-channel benefit
    let mut weight_data = Vec::with_capacity(32 * 64);
    for channel in 0..32 {
        let scale = (channel + 1) as f32 * 0.1;
        for i in 0..64 {
            weight_data.push((i as f32 * 0.01).sin() * scale);
        }
    }
    let weight = Tensor::from_slice(&weight_data, (32, 64), &device)?;

    // Per-tensor quantization
    let config_per_tensor = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer_per_tensor = Quantizer::new(config_per_tensor, device.clone());
    let quantized_pt = quantizer_per_tensor.quantize_tensor(&weight, "w1")?;
    let dequantized_pt = quantizer_per_tensor.dequantize_tensor(&quantized_pt)?;
    let error_pt = calculate_relative_error(&weight, &dequantized_pt)?;

    // Per-channel quantization
    let config_per_channel = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer_per_channel = Quantizer::new(config_per_channel, device.clone());
    let quantized_pc = quantizer_per_channel.quantize_tensor(&weight, "w1")?;
    let dequantized_pc =
        quantizer_per_channel.dequantize_tensor_per_channel(&quantized_pc, "w1")?;
    let error_pc = calculate_relative_error(&weight, &dequantized_pc)?;

    println!(
        "Variable-scale weights - Per-tensor: {:.4}%, Per-channel: {:.4}%",
        error_pt * 100.0,
        error_pc * 100.0
    );

    // Per-channel should be significantly better for variable-scale weights
    let improvement_ratio = error_pt / error_pc;
    assert!(
        improvement_ratio > 1.0,
        "Per-channel should be better than per-tensor (improvement ratio: {:.2}x)",
        improvement_ratio
    );

    // Per-channel should achieve target <1.5% error
    assert!(
        error_pc < 0.015,
        "Per-channel error {:.4}% should be < 1.5%",
        error_pc * 100.0
    );

    Ok(())
}

/// Test 5: Matmul with per-channel dequantized weights
#[test]
fn test_per_channel_matmul_integration() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Create weight [64, 128]
    let weight_data: Vec<f32> = (0..64 * 128)
        .map(|i| ((i as f32) * 0.01).sin() * 0.2)
        .collect();
    let weight = Tensor::from_slice(&weight_data, (64, 128), &device)?;

    // Create input [2, 128] (batch_size=2, in_features=128)
    let input_data: Vec<f32> = (0..2 * 128).map(|i| (i as f32) * 0.1).collect();
    let input = Tensor::from_slice(&input_data, (2, 128), &device)?;

    // F32 matmul (ground truth)
    let output_f32 = input.matmul(&weight.t()?)?;

    // Per-channel quantization
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());

    // Quantize weight
    let quantized_weight = quantizer.quantize_tensor(&weight, "weight")?;

    // Dequantize for inference
    let dequantized_weight =
        quantizer.dequantize_tensor_per_channel(&quantized_weight, "weight")?;

    // INT8 matmul (with per-channel dequantization)
    let output_int8 = input.matmul(&dequantized_weight.t()?)?;

    // Calculate output error
    let error = calculate_relative_error(&output_f32, &output_int8)?;

    println!("Matmul output error (per-channel): {:.4}%", error * 100.0);

    // Output error should be low
    assert!(
        error < 0.02,
        "Matmul output error {:.4}% should be < 2.0%",
        error * 100.0
    );

    Ok(())
}

/// Test 6: Per-channel quantization rejects non-2D tensors
#[test]
fn test_per_channel_rejects_non_2d() -> Result<(), MLError> {
    let device = Device::Cpu;

    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());

    // Test 1D tensor
    let tensor_1d = Tensor::zeros((128,), DType::F32, &device)?;
    let result_1d = quantizer.quantize_tensor_per_channel(&tensor_1d, "test");
    assert!(
        result_1d.is_err(),
        "1D tensor should be rejected for per-channel quantization"
    );

    // Test 3D tensor
    let tensor_3d = Tensor::zeros((2, 64, 128), DType::F32, &device)?;
    let result_3d = quantizer.quantize_tensor_per_channel(&tensor_3d, "test");
    assert!(
        result_3d.is_err(),
        "3D tensor should be rejected for per-channel quantization"
    );

    Ok(())
}

/// Helper: Calculate relative error between two tensors
fn calculate_relative_error(original: &Tensor, reconstructed: &Tensor) -> Result<f32, MLError> {
    // Flatten both tensors
    let orig_vec = original.flatten_all()?.to_vec1::<f32>()?;
    let recon_vec = reconstructed.flatten_all()?.to_vec1::<f32>()?;

    assert_eq!(orig_vec.len(), recon_vec.len());

    // Calculate mean absolute error
    let mae: f32 = orig_vec
        .iter()
        .zip(recon_vec.iter())
        .map(|(o, r)| (o - r).abs())
        .sum::<f32>()
        / orig_vec.len() as f32;

    // Calculate mean of original values
    let orig_mean = orig_vec.iter().sum::<f32>().abs() / orig_vec.len() as f32;

    // Relative error
    let relative_error = if orig_mean > 1e-8 {
        mae / orig_mean
    } else {
        mae // If original is near zero, use absolute error
    };

    Ok(relative_error)
}