//! TFT Temporal Self-Attention INT8 Quantization Tests
//!
//! Test-driven development for INT8 quantization of TFT attention mechanism.
//! Validates:
//! - Per-channel INT8 quantization of Q/K/V projection weights
//! - Attention score validity (no NaN/Inf)
//! - Causal masking preservation after quantization
//! - Accuracy loss <3% (stricter than other components)
//! - Memory reduction 70-80%

use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;
use ml::memory_optimization::quantization::{QuantizationConfig, QuantizationType, Quantizer};
use ml::tft::quantized_attention::QuantizedTemporalAttention;
use ml::tft::temporal_attention::TemporalSelfAttention;
use ml::MLError;

/// Test 1: Quantize Q/K/V projection weights with per-channel INT8
#[test]
fn test_quantize_qkv_projection_weights_per_channel() -> Result<(), MLError> {
    let device = Device::Cpu;
    let hidden_dim = 256;
    let num_heads = 4;

    // Create original attention module
    let vs = VarBuilder::zeros(DType::F32, &device);
    let original_attention =
        TemporalSelfAttention::new(hidden_dim, num_heads, 0.1, false, vs.pp("attention"))?;

    // Create quantization config with per-channel INT8
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(100),
    };

    // Create quantized attention module
    let quantized_attention =
        QuantizedTemporalAttention::from_attention(&original_attention, config)?;

    // Verify quantization parameters are stored per head
    let qkv_params = quantized_attention.get_quantization_params();
    assert_eq!(
        qkv_params.len(),
        num_heads * 3,
        "Should have params for Q/K/V per head"
    );

    // Verify each parameter has scale and zero_point
    for (name, params) in qkv_params.iter() {
        assert!(params.scale > 0.0, "Scale must be positive for {}", name);
        assert!(
            params.zero_point >= -128 && params.zero_point <= 127,
            "Zero point must be valid INT8 for {}",
            name
        );
    }

    println!("Test 1: Per-channel quantization PASSED");
    println!("  Quantized tensors: {}", qkv_params.len());
    println!(
        "  Example scale: {:.6}",
        qkv_params.values().next().unwrap().scale
    );

    Ok(())
}

/// Test 2: Attention scores remain valid after quantization (no NaN/Inf)
#[test]
fn test_attention_scores_validity_after_quantization() -> Result<(), MLError> {
    let device = Device::Cpu;
    let hidden_dim = 256;
    let num_heads = 4;
    let batch_size = 2;
    let seq_len = 10;

    // Create original attention
    let vs = VarBuilder::zeros(DType::F32, &device);
    let original_attention =
        TemporalSelfAttention::new(hidden_dim, num_heads, 0.1, false, vs.pp("attention"))?;

    // Create quantized attention
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(100),
    };
    let mut quantized_attention =
        QuantizedTemporalAttention::from_attention(&original_attention, config)?;

    // Create test input
    let input_data = vec![0.1f32; batch_size * seq_len * hidden_dim];
    let input = Tensor::from_slice(&input_data, (batch_size, seq_len, hidden_dim), &device)?;

    // Forward pass with quantized attention
    let output = quantized_attention.forward(&input, true)?;

    // Verify output shape
    let (out_batch, out_seq, out_dim) = output.dims3()?;
    assert_eq!(out_batch, batch_size);
    assert_eq!(out_seq, seq_len);
    assert_eq!(out_dim, hidden_dim);

    // Verify no NaN or Inf in output
    let output_vec = output.flatten_all()?.to_vec1::<f32>()?;
    let has_nan = output_vec.iter().any(|x| x.is_nan());
    let has_inf = output_vec.iter().any(|x| x.is_infinite());

    assert!(!has_nan, "Output contains NaN values");
    assert!(!has_inf, "Output contains Inf values");

    // Get attention scores and verify
    let attention_scores = quantized_attention.get_attention_scores()?;
    let scores_vec = attention_scores.flatten_all()?.to_vec1::<f32>()?;
    let scores_nan = scores_vec.iter().any(|x| x.is_nan());
    let scores_inf = scores_vec.iter().any(|x| x.is_infinite());

    assert!(!scores_nan, "Attention scores contain NaN values");
    assert!(!scores_inf, "Attention scores contain Inf values");

    // Verify attention scores sum to 1 (softmax property)
    let scores_shape = attention_scores.shape();
    let last_dim = scores_shape.dims().len() - 1;
    let sum = attention_scores.sum(last_dim)?;
    let sum_vec = sum.flatten_all()?.to_vec1::<f32>()?;
    for &s in sum_vec.iter() {
        assert!(
            (s - 1.0).abs() < 0.01,
            "Attention scores should sum to 1, got {}",
            s
        );
    }

    println!("Test 2: Attention score validity PASSED");
    println!("  Output shape: [{}, {}, {}]", out_batch, out_seq, out_dim);
    println!("  No NaN/Inf detected");
    println!("  Attention scores properly normalized");

    Ok(())
}

/// Test 3: Causal masking preserved after quantization
#[test]
fn test_causal_masking_preservation() -> Result<(), MLError> {
    let device = Device::Cpu;
    let hidden_dim = 256;
    let num_heads = 4;
    let batch_size = 1;
    let seq_len = 8;

    // Create original attention
    let vs = VarBuilder::zeros(DType::F32, &device);
    let original_attention =
        TemporalSelfAttention::new(hidden_dim, num_heads, 0.1, false, vs.pp("attention"))?;

    // Create quantized attention
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(100),
    };
    let mut quantized_attention =
        QuantizedTemporalAttention::from_attention(&original_attention, config)?;

    // Create test input with distinct values per position
    let mut input_data = Vec::new();
    for i in 0..seq_len {
        for _ in 0..hidden_dim {
            input_data.push((i as f32 + 1.0) * 0.1);
        }
    }
    let input = Tensor::from_slice(&input_data, (batch_size, seq_len, hidden_dim), &device)?;

    // Forward pass with causal masking
    let _output = quantized_attention.forward(&input, true)?;

    // Get attention scores [batch, num_heads, seq_len, seq_len]
    let attention_scores = quantized_attention.get_attention_scores()?;
    let (_, _, score_rows, score_cols) = attention_scores.dims4()?;
    assert_eq!(score_rows, seq_len);
    assert_eq!(score_cols, seq_len);

    // Extract attention scores for first head
    let head_0_scores = attention_scores.i((0, 0))?; // [seq_len, seq_len]
    let scores_2d = head_0_scores.to_vec2::<f32>()?;

    // Verify causal mask: upper triangular should be zero (or very small)
    for i in 0..seq_len {
        for j in 0..seq_len {
            if j > i {
                // Future positions should have zero attention
                assert!(
                    scores_2d[i][j] < 0.01,
                    "Causal mask violated: position ({}, {}) has attention score {:.4}",
                    i,
                    j,
                    scores_2d[i][j]
                );
            } else {
                // Past positions should have non-zero attention
                assert!(
                    scores_2d[i][j] > 0.0,
                    "Past position ({}, {}) should have attention, got {:.4}",
                    i,
                    j,
                    scores_2d[i][j]
                );
            }
        }
    }

    println!("Test 3: Causal masking preservation PASSED");
    println!("  Sequence length: {}", seq_len);
    println!("  Upper triangular (future): all zeros");
    println!("  Lower triangular (past): non-zero attention");

    Ok(())
}

/// Test 4: Accuracy loss <3% compared to FP32 (stricter than other components)
#[test]
fn test_accuracy_loss_under_3_percent() -> Result<(), MLError> {
    let device = Device::Cpu;
    let hidden_dim = 256;
    let num_heads = 4;
    let batch_size = 4;
    let seq_len = 16;

    // Create original FP32 attention
    let vs = VarBuilder::zeros(DType::F32, &device);
    let original_attention =
        TemporalSelfAttention::new(hidden_dim, num_heads, 0.1, false, vs.pp("attention"))?;

    // Create quantized INT8 attention
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(100),
    };
    let mut quantized_attention =
        QuantizedTemporalAttention::from_attention(&original_attention, config)?;

    // Generate test data
    let mut input_data = Vec::new();
    for i in 0..(batch_size * seq_len * hidden_dim) {
        let val = (i as f32 * 0.01).sin() * 0.5; // Range: [-0.5, 0.5]
        input_data.push(val);
    }
    let input = Tensor::from_slice(&input_data, (batch_size, seq_len, hidden_dim), &device)?;

    // Forward pass - original FP32
    // Note: We can't use original_attention directly because it would need to be mutable
    // So we'll compare against a cloned quantized attention with FP32 precision
    let config_fp32 = QuantizationConfig {
        quant_type: QuantizationType::None,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut fp32_attention =
        QuantizedTemporalAttention::from_attention(&original_attention, config_fp32)?;
    let output_fp32 = fp32_attention.forward(&input, true)?;

    // Forward pass - quantized INT8
    let output_int8 = quantized_attention.forward(&input, true)?;

    // Compute element-wise absolute difference
    let diff = (&output_fp32 - &output_int8)?.abs()?;
    let diff_vec = diff.flatten_all()?.to_vec1::<f32>()?;
    let max_diff = diff_vec.iter().cloned().fold(0.0f32, f32::max);

    // Compute relative error
    let fp32_vec = output_fp32.flatten_all()?.to_vec1::<f32>()?;
    let fp32_norm: f32 = fp32_vec.iter().map(|x| x * x).sum::<f32>().sqrt();

    let relative_error = if fp32_norm > 0.0 {
        (diff_vec.iter().map(|x| x * x).sum::<f32>().sqrt() / fp32_norm) * 100.0
    } else {
        0.0
    };

    println!("Test 4: Accuracy loss PASSED");
    println!("  Relative error: {:.4}%", relative_error);
    println!("  Max absolute difference: {:.6}", max_diff);
    println!("  Target: <3%");

    assert!(
        relative_error < 3.0,
        "Accuracy loss {:.4}% exceeds 3% threshold",
        relative_error
    );

    Ok(())
}

/// Test 5: Memory reduction 70-80%
#[test]
fn test_memory_reduction_70_to_80_percent() -> Result<(), MLError> {
    let device = Device::Cpu;
    let hidden_dim = 256;
    let num_heads = 4;

    // Create original attention
    let vs = VarBuilder::zeros(DType::F32, &device);
    let original_attention =
        TemporalSelfAttention::new(hidden_dim, num_heads, 0.1, false, vs.pp("attention"))?;

    // Calculate original FP32 memory size
    let head_dim = hidden_dim / num_heads;
    let qkv_size_per_head = hidden_dim * head_dim; // Input dim × output dim
    let total_qkv_params = num_heads * 3 * qkv_size_per_head; // Q, K, V for each head
    let fp32_bytes = total_qkv_params * 4; // 4 bytes per float32

    // Create quantized attention
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(100),
    };
    let quantized_attention =
        QuantizedTemporalAttention::from_attention(&original_attention, config)?;

    // Get quantized memory size
    let int8_bytes = quantized_attention.memory_bytes();

    // Calculate reduction
    let reduction_percent = ((fp32_bytes - int8_bytes) as f64 / fp32_bytes as f64) * 100.0;

    println!("Test 5: Memory reduction PASSED");
    println!("  FP32 size: {} bytes ({:.2} MB)", fp32_bytes, fp32_bytes as f64 / 1_048_576.0);
    println!("  INT8 size: {} bytes ({:.2} MB)", int8_bytes, int8_bytes as f64 / 1_048_576.0);
    println!("  Reduction: {:.2}%", reduction_percent);
    println!("  Target: 70-80%");

    assert!(
        reduction_percent >= 70.0 && reduction_percent <= 80.0,
        "Memory reduction {:.2}% not in 70-80% range",
        reduction_percent
    );

    Ok(())
}

/// Test 6: Quantization with different batch sizes
#[test]
fn test_quantization_with_various_batch_sizes() -> Result<(), MLError> {
    let device = Device::Cpu;
    let hidden_dim = 256;
    let num_heads = 4;
    let seq_len = 10;

    // Create quantized attention
    let vs = VarBuilder::zeros(DType::F32, &device);
    let original_attention =
        TemporalSelfAttention::new(hidden_dim, num_heads, 0.1, false, vs.pp("attention"))?;

    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(100),
    };
    let mut quantized_attention =
        QuantizedTemporalAttention::from_attention(&original_attention, config)?;

    // Test with different batch sizes
    for batch_size in [1, 4, 16, 32] {
        let input_data = vec![0.1f32; batch_size * seq_len * hidden_dim];
        let input = Tensor::from_slice(&input_data, (batch_size, seq_len, hidden_dim), &device)?;

        let output = quantized_attention.forward(&input, true)?;
        let (out_batch, out_seq, out_dim) = output.dims3()?;

        assert_eq!(out_batch, batch_size);
        assert_eq!(out_seq, seq_len);
        assert_eq!(out_dim, hidden_dim);

        // Verify no NaN/Inf
        let output_vec = output.flatten_all()?.to_vec1::<f32>()?;
        let has_nan = output_vec.iter().any(|x| x.is_nan());
        let has_inf = output_vec.iter().any(|x| x.is_infinite());

        assert!(
            !has_nan && !has_inf,
            "Batch size {} produced NaN/Inf",
            batch_size
        );
    }

    println!("Test 6: Various batch sizes PASSED");
    println!("  Tested batch sizes: [1, 4, 16, 32]");
    println!("  All outputs valid");

    Ok(())
}

/// Test 7: Dequantization accuracy
#[test]
fn test_dequantization_accuracy() -> Result<(), MLError> {
    let device = Device::Cpu;
    let hidden_dim = 256;
    let num_heads = 4;

    // Create original attention
    let vs = VarBuilder::zeros(DType::F32, &device);
    let original_attention =
        TemporalSelfAttention::new(hidden_dim, num_heads, 0.1, false, vs.pp("attention"))?;

    // Create quantized attention
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(100),
    };
    let quantized_attention =
        QuantizedTemporalAttention::from_attention(&original_attention, config)?;

    // Test dequantization of each Q/K/V projection
    let qkv_params = quantized_attention.get_quantization_params();

    for (name, params) in qkv_params.iter() {
        // Verify scale is reasonable (not too small or too large)
        assert!(
            params.scale > 1e-6 && params.scale < 1e6,
            "Scale {} out of range for {}",
            params.scale,
            name
        );

        // Verify min/max range is captured
        assert!(
            params.min_val <= params.max_val,
            "Invalid min/max range for {}",
            name
        );
    }

    println!("Test 7: Dequantization accuracy PASSED");
    println!("  All quantization parameters valid");
    println!("  Scale ranges verified");

    Ok(())
}