#![cfg(feature = "disabled_tests")]

use anyhow::Result;
use rtx_compress::kv_cache::{CompressedKVCache, CompressionMethod, KVCacheConfig};
use rtx_tensor::{Device, Shape, Tensor};

#[test]
fn test_kv_cache_compression_ratio() -> Result<()> {
    let device = Device::try_default()?;

    // Create sample KV cache: [batch, num_heads, seq_len, head_dim]
    let keys = Tensor::randn(&[2, 12, 1024, 64], &device)?;
    let values = Tensor::randn(&[2, 12, 1024, 64], &device)?;

    let config = KVCacheConfig {
        compression_method: CompressionMethod::ProductQuantization {
            num_subquantizers: 8,
            codebook_size: 256,
            use_opq: true,
        },
        compression_ratio_target: 4.0,
        quality_threshold: 0.95, // 95% quality retention
        max_cache_size_mb: 1024,
        enable_sliding_window: true,
        window_size: 2048,
        enable_attention_scoring: true,
        prefetch_batch_size: 32,
        enable_auto_tuning: false,
    };

    let mut cache = CompressedKVCache::new(config)?;

    // Insert KV pairs
    cache.insert(0, &keys, &values)?;

    // Verify compression metrics
    let stats = cache.compression_stats();
    assert!(
        stats.compression_ratio >= 3.5,
        "Compression ratio should be >= 3.5x, got {}",
        stats.compression_ratio
    );
    assert!(
        stats.quality_score >= 0.90,
        "Quality should be >= 90%, got {}",
        stats.quality_score
    );

    // Test retrieval
    let (retrieved_keys, retrieved_values) = cache.get(0, 0, 1024)?;

    // Verify shapes match
    assert_eq!(retrieved_keys.shape(), keys.shape());
    assert_eq!(retrieved_values.shape(), values.shape());

    // Verify quality
    let key_mse = (&keys - &retrieved_keys)?
        .pow(2.0)?
        .mean(&[], false)?
        .to_scalar::<f32>()?;
    let value_mse = (&values - &retrieved_values)?
        .pow(2.0)?
        .mean(&[], false)?
        .to_scalar::<f32>()?;

    assert!(key_mse < 0.01, "Key MSE should be < 0.01, got {}", key_mse);
    assert!(
        value_mse < 0.01,
        "Value MSE should be < 0.01, got {}",
        value_mse
    );

    Ok(())
}

#[test]
fn test_kv_cache_incremental_updates() -> Result<()> {
    let device = Device::try_default()?;

    let config = KVCacheConfig {
        compression_method: CompressionMethod::VectorQuantization {
            codebook_size: 1024,
            update_frequency: 100,
        },
        compression_ratio_target: 3.0,
        quality_threshold: 0.92,
        max_cache_size_mb: 1024,
        enable_sliding_window: false,
        window_size: 1024,
        enable_attention_scoring: false,
        prefetch_batch_size: 16,
        enable_auto_tuning: false,
    };

    let mut cache = CompressedKVCache::new(config)?;

    // Start with small sequence
    let keys1 = Tensor::randn(&[1, 8, 64, 32], &device)?;
    let values1 = Tensor::randn(&[1, 8, 64, 32], &device)?;

    cache.insert(0, &keys1, &values1)?;

    // Add more tokens incrementally
    for i in 1..10 {
        let new_keys = Tensor::randn(&[1, 8, 1, 32], &device)?;
        let new_values = Tensor::randn(&[1, 8, 1, 32], &device)?;

        cache.append(0, &new_keys, &new_values)?;

        // Verify total sequence length
        let (retrieved_k, retrieved_v) = cache.get(0, 0, 64 + i)?;
        assert_eq!(retrieved_k.shape()[2], 64 + i);
        assert_eq!(retrieved_v.shape()[2], 64 + i);
    }

    // Verify final compression stats
    let stats = cache.compression_stats();
    assert!(stats.compression_ratio >= 2.5);

    Ok(())
}

#[test]
fn test_kv_cache_memory_management() -> Result<()> {
    let device = Device::try_default()?;

    let config = KVCacheConfig {
        compression_method: CompressionMethod::MixedPrecision {
            fp16_layers: vec![0, 1, 2],
            int8_layers: vec![3, 4, 5],
            int4_layers: vec![6, 7],
        },
        compression_ratio_target: 6.0,
        quality_threshold: 0.88,
        max_cache_size_mb: 1024,
        enable_sliding_window: false,
        window_size: 1024,
        enable_attention_scoring: false,
        prefetch_batch_size: 16,
        enable_auto_tuning: false,
    };

    let mut cache = CompressedKVCache::new(config)?;
    cache.set_memory_limit(1_000_000); // 1MB limit

    // Fill cache beyond memory limit
    for seq_id in 0..100 {
        let keys = Tensor::randn(&[1, 4, 128, 16], &device)?;
        let values = Tensor::randn(&[1, 4, 128, 16], &device)?;

        cache.insert(seq_id, &keys, &values)?;
    }

    // Verify memory usage is within limits
    let memory_usage = cache.memory_usage();
    assert!(
        memory_usage <= 1_100_000, // Allow 10% overhead
        "Memory usage {} exceeds limit",
        memory_usage
    );

    // Verify LRU eviction works
    assert!(cache.contains_sequence(99)); // Most recent should be present
    assert!(!cache.contains_sequence(0)); // Oldest should be evicted

    Ok(())
}

#[test]
fn test_kv_cache_batch_operations() -> Result<()> {
    let device = Device::try_default()?;

    let config = KVCacheConfig {
        compression_method: CompressionMethod::ProductQuantization {
            num_subquantizers: 4,
            codebook_size: 128,
            use_opq: false,
        },
        compression_ratio_target: 4.0,
        quality_threshold: 0.93,
        max_cache_size_mb: 512,
        enable_sliding_window: false,
        window_size: 512,
        enable_attention_scoring: false,
        prefetch_batch_size: 8,
        enable_auto_tuning: false,
    };

    let mut cache = CompressedKVCache::new(config)?;

    // Insert batch of sequences
    let batch_keys = Tensor::randn(&[8, 6, 256, 24], &device)?;
    let batch_values = Tensor::randn(&[8, 6, 256, 24], &device)?;

    cache.insert_batch(&[0, 1, 2, 3, 4, 5, 6, 7], &batch_keys, &batch_values)?;

    // Retrieve batch
    let (retrieved_keys, retrieved_values) = cache.get_batch(&[0, 2, 4, 6], 0, 256)?;

    assert_eq!(retrieved_keys.shape().dims(), &[4, 6, 256, 24]);
    assert_eq!(retrieved_values.shape().dims(), &[4, 6, 256, 24]);

    Ok(())
}

#[test]
fn test_kv_cache_attention_pattern_optimization() -> Result<()> {
    let device = Device::try_default()?;

    let config = KVCacheConfig {
        compression_method: CompressionMethod::AdaptiveQuantization {
            base_bits: 8,
            attention_threshold: 0.1,
            importance_decay: 0.95,
        },
        compression_ratio_target: 5.0,
        quality_threshold: 0.90,
        max_cache_size_mb: 1024,
        enable_sliding_window: false,
        window_size: 1024,
        enable_attention_scoring: true,
        prefetch_batch_size: 16,
        enable_auto_tuning: false,
    };

    let mut cache = CompressedKVCache::new(config)?;

    let keys = Tensor::randn(&[1, 4, 512, 32], &device)?;
    let values = Tensor::randn(&[1, 4, 512, 32], &device)?;

    // Simulate attention patterns (higher attention = more important)
    let attention_weights = Tensor::randn(&[1, 4, 512], &device)?.softmax(-1)?;

    cache.insert_with_attention(0, &keys, &values, &attention_weights)?;

    // Important tokens should have better quality
    let (retrieved_k, retrieved_v) = cache.get(0, 0, 512)?;

    // Compute per-token reconstruction error
    let key_errors = (keys - retrieved_k)?.pow(2.0)?.mean(&[-1])?;
    let attention_flat = attention_weights.mean(&[1])?; // Average across heads

    // Verify inverse correlation between attention and error
    let correlation = compute_correlation(&attention_flat, &key_errors)?;
    assert!(
        correlation < -0.1,
        "Higher attention should correlate with lower error"
    );

    Ok(())
}

// Helper function to compute correlation
fn compute_correlation(x: &Tensor, y: &Tensor) -> Result<f32> {
    let x_mean = x.mean(&[], false)?;
    let y_mean = y.mean(&[], false)?;

    let x_centered = (x - &x_mean)?;
    let y_centered = (y - &y_mean)?;

    let numerator = (&x_centered * &y_centered)?.sum(Some(0))?;
    let x_var = x_centered.pow(2.0)?.sum(Some(0))?;
    let y_var = y_centered.pow(2.0)?.sum(Some(0))?;

    let correlation = &numerator / &(&x_var * &y_var)?.pow(0.5)?;
    correlation.to_scalar::<f32>()
}

#[test]
fn test_kv_cache_persistence() -> Result<()> {
    let device = Device::try_default()?;

    let config = KVCacheConfig {
        compression_method: CompressionMethod::ProductQuantization {
            num_subquantizers: 2,
            codebook_size: 64,
            use_opq: false,
        },
        compression_ratio_target: 3.0,
        quality_threshold: 0.90,
        max_cache_size_mb: 256,
        enable_sliding_window: false,
        window_size: 512,
        enable_attention_scoring: false,
        prefetch_batch_size: 4,
        enable_auto_tuning: false,
    };

    let mut cache = CompressedKVCache::new(config.clone())?;

    let keys = Tensor::randn(&[1, 2, 128, 16], &device)?;
    let values = Tensor::randn(&[1, 2, 128, 16], &device)?;

    cache.insert(42, &keys, &values)?;

    // Save cache
    let saved_data = cache.save_to_bytes()?;

    // Load into new cache
    let mut cache2 = CompressedKVCache::new(config)?;
    cache2.load_from_bytes(&saved_data)?;

    // Verify data is preserved
    let (retrieved_k, retrieved_v) = cache2.get(42, 0, 128)?;

    let key_diff = (&keys - &retrieved_k)?.abs()?.max()?.to_scalar::<f32>()?;
    let value_diff = (&values - &retrieved_v)?.abs()?.max()?.to_scalar::<f32>()?;

    assert!(key_diff < 0.1, "Persisted keys should match original");
    assert!(value_diff < 0.1, "Persisted values should match original");

    Ok(())
}