//! TFT Gated Residual Network INT8 Quantization Tests
//!
//! Test-driven development for GRN INT8 quantization with residual connections.
//! Target: 500MB → 125MB (75% reduction) with <5% accuracy loss.

use candle_core::{DType, Device, Tensor};
use candle_nn::{VarBuilder, VarMap};
use std::sync::Arc;

use ml::tft::gated_residual::GatedResidualNetwork;
use ml::memory_optimization::quantization::{QuantizationConfig, QuantizationType, Quantizer};
use ml::tft::quantized_grn::QuantizedGatedResidualNetwork;
use ml::MLError;

/// Test 1: Quantize GRN linear layers to INT8
#[test]
fn test_quantize_grn_linear_layers() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create original GRN
    let grn = GatedResidualNetwork::new(128, 128, vs.pp("grn"))?;

    // Quantization config
    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(100),
    };

    let quantizer = Quantizer::new(quant_config, device.clone());

    // Create quantized GRN
    let quantized_grn = QuantizedGatedResidualNetwork::from_grn(&grn, quantizer)?;

    // Verify quantization occurred
    assert_eq!(quantized_grn.quant_type(), QuantizationType::Int8);
    assert!(quantized_grn.quantized_linear1.is_some());
    assert!(quantized_grn.quantized_linear2.is_some());
    assert!(quantized_grn.quantized_glu_weights.0.is_some());
    assert!(quantized_grn.quantized_glu_weights.1.is_some());

    Ok(())
}

/// Test 2: Skip connection accuracy maintained in F32
#[test]
fn test_skip_connection_accuracy() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create GRN with dimension mismatch (requires skip projection)
    let grn = GatedResidualNetwork::new(64, 128, vs.pp("grn"))?;

    // Create test input
    let input_data = vec![1.0f32; 128]; // batch=2, dim=64
    let input = Tensor::from_slice(&input_data, (2, 64), &device)?;

    // Original forward pass
    let original_output = grn.forward(&input, None)?;

    // Quantize GRN
    let quant_config = QuantizationConfig::default();
    let quantizer = Quantizer::new(quant_config, device.clone());
    let quantized_grn = QuantizedGatedResidualNetwork::from_grn(&grn, quantizer)?;

    // Quantized forward pass
    let quantized_output = quantized_grn.forward(&input, None)?;

    // Calculate difference
    let diff = (&original_output - &quantized_output)?;
    let diff_vec = diff.flatten_all()?.to_vec1::<f32>()?;
    let mae = diff_vec.iter().map(|x| x.abs()).sum::<f32>() / diff_vec.len() as f32;

    // Skip connection should be high precision (kept in F32)
    // MAE should be < 0.1 (10% of typical value range)
    println!("Skip connection MAE: {:.6}", mae);
    assert!(mae < 0.1, "Skip connection error too high: {}", mae);

    Ok(())
}

/// Test 3: Gating mechanism works with INT8
#[test]
fn test_gating_mechanism_int8() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create GRN
    let grn = GatedResidualNetwork::new(128, 128, vs.pp("grn"))?;

    // Create test input
    let input_data = vec![0.5f32; 256]; // batch=2, dim=128
    let input = Tensor::from_slice(&input_data, (2, 128), &device)?;

    // Original GLU output
    let original_output = grn.forward(&input, None)?;

    // Quantize
    let quant_config = QuantizationConfig::default();
    let quantizer = Quantizer::new(quant_config, device.clone());
    let quantized_grn = QuantizedGatedResidualNetwork::from_grn(&grn, quantizer)?;

    // Quantized GLU output
    let quantized_output = quantized_grn.forward(&input, None)?;

    // Check gating still produces valid outputs (not NaN, not Inf)
    let output_vec = quantized_output.flatten_all()?.to_vec1::<f32>()?;
    assert!(output_vec.iter().all(|x| x.is_finite()), "Gating produced invalid values");

    // Check gating behavior preserved (output should be in reasonable range)
    let mean = output_vec.iter().sum::<f32>() / output_vec.len() as f32;
    println!("Quantized gating output mean: {:.6}", mean);
    assert!(mean.abs() < 10.0, "Gating output out of range");

    Ok(())
}

/// Test 4: Accuracy loss < 5%
#[test]
fn test_accuracy_loss_under_5_percent() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create GRN
    let grn = GatedResidualNetwork::new(128, 128, vs.pp("grn"))?;

    // Create diverse test inputs
    let num_samples = 100;
    let mut total_relative_error = 0.0;

    for i in 0..num_samples {
        // Generate varying inputs
        let scale = 1.0 + (i as f32) * 0.01;
        let input_data = vec![scale; 256]; // batch=2, dim=128
        let input = Tensor::from_slice(&input_data, (2, 128), &device)?;

        // Original output
        let original = grn.forward(&input, None)?;
        let original_vec = original.flatten_all()?.to_vec1::<f32>()?;

        // Quantized output
        let quant_config = QuantizationConfig::default();
        let quantizer = Quantizer::new(quant_config, device.clone());
        let quantized_grn = QuantizedGatedResidualNetwork::from_grn(&grn, quantizer)?;
        let quantized = quantized_grn.forward(&input, None)?;
        let quantized_vec = quantized.flatten_all()?.to_vec1::<f32>()?;

        // Calculate relative error
        let mut sample_error = 0.0;
        for (orig, quant) in original_vec.iter().zip(quantized_vec.iter()) {
            let relative_err = (orig - quant).abs() / (orig.abs() + 1e-8);
            sample_error += relative_err;
        }
        sample_error /= original_vec.len() as f32;
        total_relative_error += sample_error;
    }

    let avg_relative_error = total_relative_error / num_samples as f32;
    println!("Average relative error: {:.4}%", avg_relative_error * 100.0);

    // Assert < 5% accuracy loss
    assert!(
        avg_relative_error < 0.05,
        "Accuracy loss {:.2}% exceeds 5% threshold",
        avg_relative_error * 100.0
    );

    Ok(())
}

/// Test 5: Memory reduction 70-80%
#[test]
fn test_memory_reduction_70_to_80_percent() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create GRN with known size
    let input_dim = 512;
    let output_dim = 512;
    let grn = GatedResidualNetwork::new(input_dim, output_dim, vs.pp("grn"))?;

    // Calculate original memory footprint
    // linear1: 512 × 512 × 4 bytes = 1,048,576 bytes
    // linear2: 512 × 512 × 4 bytes = 1,048,576 bytes
    // glu.linear: 512 × 512 × 4 bytes = 1,048,576 bytes
    // glu.gate: 512 × 512 × 4 bytes = 1,048,576 bytes
    // skip_projection: None (same dims)
    // Total: ~4.0 MB
    let original_memory_mb = 4.0;

    // Quantize
    let quant_config = QuantizationConfig::default();
    let quantizer = Quantizer::new(quant_config, device.clone());
    let quantized_grn = QuantizedGatedResidualNetwork::from_grn(&grn, quantizer)?;

    // Calculate quantized memory footprint
    let quantized_memory_mb = quantized_grn.memory_footprint_mb();

    // Calculate reduction percentage
    let reduction_percent = (1.0 - quantized_memory_mb / original_memory_mb) * 100.0;
    println!("Memory reduction: {:.1}% ({:.2} MB → {:.2} MB)",
             reduction_percent, original_memory_mb, quantized_memory_mb);

    // Assert 70-80% reduction (INT8 should give ~75%)
    assert!(
        reduction_percent >= 70.0 && reduction_percent <= 80.0,
        "Memory reduction {:.1}% not in 70-80% range",
        reduction_percent
    );

    Ok(())
}

/// Test 6: Quantized GRN forward pass with context
#[test]
fn test_quantized_forward_with_context() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create GRN
    let grn = GatedResidualNetwork::new(128, 128, vs.pp("grn"))?;

    // Create test input and context
    let input_data = vec![1.0f32; 256]; // batch=2, dim=128
    let input = Tensor::from_slice(&input_data, (2, 128), &device)?;

    let context_data = vec![0.5f32; 256]; // batch=2, dim=128
    let context = Tensor::from_slice(&context_data, (2, 128), &device)?;

    // Original output with context
    let original_output = grn.forward(&input, Some(&context))?;

    // Quantize
    let quant_config = QuantizationConfig::default();
    let quantizer = Quantizer::new(quant_config, device.clone());
    let quantized_grn = QuantizedGatedResidualNetwork::from_grn(&grn, quantizer)?;

    // Quantized output with context
    let quantized_output = quantized_grn.forward(&input, Some(&context))?;

    // Verify shapes match
    assert_eq!(original_output.dims(), quantized_output.dims());

    // Calculate accuracy
    let diff = (&original_output - &quantized_output)?;
    let diff_vec = diff.flatten_all()?.to_vec1::<f32>()?;
    let mae = diff_vec.iter().map(|x| x.abs()).sum::<f32>() / diff_vec.len() as f32;

    println!("Context forward MAE: {:.6}", mae);
    assert!(mae < 0.2, "Context forward error too high: {}", mae);

    Ok(())
}