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

use anyhow::Result;
use rtx_compress::{
    CompressedStorage, CompressionConfig,
    checkpoint::CheckpointCompressor,
    kv_cache::CompressedKVCache,
    quantization::{
        mixed_precision::MixedPrecisionOptimizer, product_quantization::ProductQuantizer,
    },
};
use rtx_inference::{InferenceEngine, InferenceRequest};
use rtx_tensor::{DType, Device, Tensor};
use std::collections::HashMap;

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

    // Step 1: Create a mock transformer model
    let mut model_weights = HashMap::new();
    for layer in 0..6 {
        // 6-layer transformer
        model_weights.insert(
            format!("layers.{}.attention.query", layer),
            Tensor::randn(&[512, 512], &device)?,
        );
        model_weights.insert(
            format!("layers.{}.attention.key", layer),
            Tensor::randn(&[512, 512], &device)?,
        );
        model_weights.insert(
            format!("layers.{}.attention.value", layer),
            Tensor::randn(&[512, 512], &device)?,
        );
        model_weights.insert(
            format!("layers.{}.mlp.up_proj", layer),
            Tensor::randn(&[512, 2048], &device)?,
        );
        model_weights.insert(
            format!("layers.{}.mlp.down_proj", layer),
            Tensor::randn(&[2048, 512], &device)?,
        );
    }

    // Step 2: Apply mixed precision optimization
    let precision_config = rtx_compress::quantization::mixed_precision::PrecisionConfig {
        precision_bits: vec![4, 8, 12, 16],
        sensitivity_threshold: 0.02,
        performance_weight: 0.6,
        quality_weight: 0.4,
    };

    let mut precision_optimizer = MixedPrecisionOptimizer::new(precision_config);
    let calibration_data = Tensor::randn(&[100, 512], &device)?;

    let objective = rtx_compress::quantization::mixed_precision::OptimizationObjective {
        target_compression_ratio: 4.0,
        max_quality_loss: 0.05,
        memory_constraint_mb: Some(200),
    };

    let optimal_precision =
        precision_optimizer.optimize(&model_weights, &calibration_data, objective)?;

    // Step 3: Compress model checkpoints
    let checkpoint_config = rtx_compress::checkpoint::CompressionConfig {
        format: rtx_compress::checkpoint::CheckpointFormat::Zstd,
        compression_level: 6,
        quantization_bits: 0, // Use mixed precision
        exclude_patterns: vec!["*layer_norm*".to_string()],
    };

    let mut checkpoint_compressor = CheckpointCompressor::new(checkpoint_config);

    // Apply optimal precisions
    for (layer, bits) in &optimal_precision.layer_precisions {
        checkpoint_compressor.add_quantization_rule(layer, *bits);
    }

    let compressed_model = checkpoint_compressor.save(&model_weights)?;

    // Step 4: Set up compressed KV cache for inference
    let kv_config = rtx_compress::kv_cache::KVCacheConfig {
        compression_method: rtx_compress::kv_cache::CompressionMethod::ProductQuantization {
            num_subquantizers: 8,
            codebook_size: 256,
            use_opq: true,
        },
        compression_ratio_target: 3.0,
        quality_threshold: 0.92,
        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 kv_cache = CompressedKVCache::new(kv_config)?;

    // Step 5: Verify full pipeline works
    let loaded_model = checkpoint_compressor.load(&compressed_model)?;
    assert_eq!(loaded_model.len(), model_weights.len());

    // Simulate inference with compressed KV cache
    for seq_id in 0..5 {
        let keys = Tensor::randn(&[1, 8, 256, 64], &device)?;
        let values = Tensor::randn(&[1, 8, 256, 64], &device)?;

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

    // Verify compression metrics
    let compression_stats = kv_cache.compression_stats();
    assert!(compression_stats.compression_ratio >= 2.5);

    let original_model_size = calculate_model_size(&model_weights);
    let compressed_model_size = compressed_model.len();
    let model_compression_ratio = original_model_size as f64 / compressed_model_size as f64;

    assert!(
        model_compression_ratio >= 3.0,
        "Model compression should achieve 3x ratio, got {:.2}",
        model_compression_ratio
    );

    println!("Pipeline results:");
    println!("  Model compression: {:.2}x", model_compression_ratio);
    println!(
        "  KV cache compression: {:.2}x",
        compression_stats.compression_ratio
    );
    println!(
        "  Total memory savings: ~{:.1}%",
        (1.0 - 1.0 / model_compression_ratio) * 100.0
    );

    Ok(())
}

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

    // Create compressed storage system
    let storage_config = CompressionConfig {
        default_compression_ratio: 4.0,
        quality_threshold: 0.90,
        memory_limit_mb: 1000,
        adaptive_compression: true,
    };

    let mut storage = CompressedStorage::new(storage_config);

    // Store model layers with different compression strategies
    let attention_weights = Tensor::randn(&[768, 768], &device)?;
    let embedding_weights = Tensor::randn(&[50000, 768], &device)?;
    let mlp_weights = Tensor::randn(&[768, 3072], &device)?;

    // High precision for attention (critical for quality)
    storage.store_tensor(
        "attention.query",
        &attention_weights,
        rtx_compress::CompressionLevel::Low,
    )?;

    // Aggressive compression for embeddings (large but redundant)
    storage.store_tensor(
        "embeddings.weight",
        &embedding_weights,
        rtx_compress::CompressionLevel::High,
    )?;

    // Balanced compression for MLP
    storage.store_tensor(
        "mlp.up_proj",
        &mlp_weights,
        rtx_compress::CompressionLevel::Medium,
    )?;

    // Simulate inference requests
    let batch_size = 4;
    let seq_len = 512;

    for request_id in 0..10 {
        // Load weights on demand (decompression)
        let attention_loaded = storage.load_tensor("attention.query")?;
        let embeddings_loaded = storage.load_tensor("embeddings.weight")?;
        let mlp_loaded = storage.load_tensor("mlp.up_proj")?;

        // Verify shapes are preserved
        assert_eq!(attention_loaded.shape(), attention_weights.shape());
        assert_eq!(embeddings_loaded.shape(), embedding_weights.shape());
        assert_eq!(mlp_loaded.shape(), mlp_weights.shape());

        // Simulate forward pass computation
        let input_ids = Tensor::randint(0, 50000, &[batch_size, seq_len], &device)?;
        let embeddings = embeddings_loaded.gather(&input_ids, 0)?;

        // Mock attention computation
        let queries = embeddings.matmul(&attention_loaded)?;

        // Verify computation produces reasonable results
        assert_eq!(queries.shape(), &[batch_size, seq_len, 768]);

        println!(
            "Request {}: processed batch_size={}, seq_len={}",
            request_id, batch_size, seq_len
        );
    }

    // Verify storage efficiency
    let storage_stats = storage.get_statistics();
    println!("Storage statistics:");
    println!(
        "  Total compressed size: {} MB",
        storage_stats.total_compressed_size_mb
    );
    println!(
        "  Average compression ratio: {:.2}x",
        storage_stats.average_compression_ratio
    );
    println!(
        "  Memory usage: {:.1}%",
        storage_stats.memory_utilization * 100.0
    );

    assert!(storage_stats.average_compression_ratio >= 2.0);
    assert!(storage_stats.memory_utilization <= 1.0);

    Ok(())
}

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

    // Simulate multiple model variants that share common layers
    let shared_embeddings = Tensor::randn(&[30000, 512], &device)?;
    let shared_attention = Tensor::randn(&[512, 512], &device)?;

    // Model A: Base model
    let mut model_a = HashMap::new();
    model_a.insert("embeddings".to_string(), shared_embeddings.clone());
    model_a.insert("attention".to_string(), shared_attention.clone());
    model_a.insert(
        "classifier".to_string(),
        Tensor::randn(&[512, 1000], &device)?,
    );

    // Model B: Fine-tuned variant (shares embeddings and attention)
    let mut model_b = HashMap::new();
    model_b.insert("embeddings".to_string(), shared_embeddings.clone());
    model_b.insert("attention".to_string(), shared_attention.clone());
    model_b.insert(
        "classifier".to_string(),
        Tensor::randn(&[512, 100], &device)?,
    );

    let config = rtx_compress::checkpoint::CompressionConfig {
        format: rtx_compress::checkpoint::CheckpointFormat::Zstd,
        compression_level: 5,
        quantization_bits: 8,
        exclude_patterns: vec![],
    };

    let compressor = CheckpointCompressor::new(config);

    // Save models with deduplication
    let compressed_a = compressor.save(&model_a)?;
    let compressed_b_incremental = compressor.save_with_base(&model_b, &model_a)?;

    // Verify incremental compression is much smaller
    let compressed_b_full = compressor.save(&model_b)?;

    assert!(
        compressed_b_incremental.len() < compressed_b_full.len() / 2,
        "Incremental compression should be much smaller"
    );

    // Verify both models can be loaded correctly
    let loaded_a = compressor.load(&compressed_a)?;
    let loaded_b = compressor.load_with_base(&compressed_b_incremental, &model_a)?;

    // Verify shared layers are identical
    let embeddings_diff = (&loaded_a["embeddings"] - &loaded_b["embeddings"])?
        .abs()?
        .max()?
        .to_scalar::<f32>()?;
    assert!(
        embeddings_diff < 1e-6,
        "Shared embeddings should be identical"
    );

    let attention_diff = (&loaded_a["attention"] - &loaded_b["attention"])?
        .abs()?
        .max()?
        .to_scalar::<f32>()?;
    assert!(
        attention_diff < 1e-6,
        "Shared attention should be identical"
    );

    // But classifiers should be different
    let classifier_diff = (&loaded_a["classifier"] - &loaded_b["classifier"])?
        .abs()?
        .mean(&[], false)?
        .to_scalar::<f32>()?;
    assert!(classifier_diff > 0.1, "Classifiers should be different");

    println!("Model sharing results:");
    println!("  Model A size: {} bytes", compressed_a.len());
    println!("  Model B (full): {} bytes", compressed_b_full.len());
    println!(
        "  Model B (incremental): {} bytes",
        compressed_b_incremental.len()
    );
    println!(
        "  Sharing savings: {:.1}%",
        (1.0 - compressed_b_incremental.len() as f64 / compressed_b_full.len() as f64) * 100.0
    );

    Ok(())
}

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

    let storage_config = CompressionConfig {
        default_compression_ratio: 3.0,
        quality_threshold: 0.88,
        memory_limit_mb: 500, // Constrained memory
        adaptive_compression: true,
    };

    let mut storage = CompressedStorage::new(storage_config);
    storage.enable_load_monitoring(true);

    // Store many tensors to exceed memory limit
    let mut tensors = HashMap::new();
    for i in 0..100 {
        let tensor_name = format!("layer_{}", i);
        let tensor = Tensor::randn(&[1024, 1024], &device)?;
        tensors.insert(tensor_name.clone(), tensor.clone());

        // Simulate varying access patterns
        let access_frequency = if i < 20 {
            rtx_compress::AccessPattern::VeryHigh
        } else if i < 50 {
            rtx_compress::AccessPattern::Medium
        } else {
            rtx_compress::AccessPattern::Low
        };

        storage.store_tensor_with_hint(&tensor_name, &tensor, access_frequency)?;
    }

    // Simulate high load with concurrent access
    let mut handles = vec![];
    for thread_id in 0..4 {
        let storage_clone = storage.clone();
        let handle = std::thread::spawn(move || -> Result<()> {
            for request in 0..25 {
                let tensor_idx = (thread_id * 25 + request) % 100;
                let tensor_name = format!("layer_{}", tensor_idx);

                // Load tensor (may trigger adaptive recompression)
                let loaded = storage_clone.load_tensor(&tensor_name)?;

                // Simulate computation
                let result = loaded.sum(Some(0))?;
                let scalar = result.to_vec::<f32>()?;
                assert!(scalar[0].is_finite());
            }
            Ok(())
        });
        handles.push(handle);
    }

    // Wait for all threads
    for handle in handles {
        handle.join().unwrap()?;
    }

    // Verify adaptive behavior
    let final_stats = storage.get_statistics();
    println!("Adaptive compression results:");
    println!(
        "  Memory usage: {:.1}%",
        final_stats.memory_utilization * 100.0
    );
    println!(
        "  Average compression ratio: {:.2}x",
        final_stats.average_compression_ratio
    );
    println!("  Cache hits: {:.1}%", final_stats.cache_hit_rate * 100.0);
    println!("  Evictions: {}", final_stats.eviction_count);

    // Should stay within memory limits
    assert!(
        final_stats.memory_utilization <= 1.0,
        "Should not exceed memory limit"
    );

    // Should achieve reasonable cache performance under load
    assert!(
        final_stats.cache_hit_rate >= 0.7,
        "Should achieve reasonable cache hit rate under load"
    );

    Ok(())
}

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

    // Create test tensor with known characteristics
    let test_tensor = Tensor::randn(&[2048, 2048], &device)?;

    let compression_levels = vec![
        (
            "ultra_fast",
            rtx_compress::CompressionLevel::UltraFast,
            10.0,
        ), // 10ms target
        ("fast", rtx_compress::CompressionLevel::Fast, 50.0), // 50ms target
        ("balanced", rtx_compress::CompressionLevel::Balanced, 200.0), // 200ms target
        ("high", rtx_compress::CompressionLevel::High, 1000.0), // 1s target
    ];

    let mut results = vec![];

    for (name, level, time_budget_ms) in compression_levels {
        let config = CompressionConfig {
            default_compression_ratio: 0.0, // Auto-determine based on level
            quality_threshold: 0.0,         // Auto-determine
            memory_limit_mb: 1000,
            adaptive_compression: false,
        };

        let mut storage = CompressedStorage::new(config);

        // Measure compression time
        let start = std::time::Instant::now();
        storage.store_tensor_with_level("test", &test_tensor, level)?;
        let compression_time = start.elapsed();

        // Measure decompression time
        let start = std::time::Instant::now();
        let decompressed = storage.load_tensor("test")?;
        let decompression_time = start.elapsed();

        // Measure quality
        let mse = (&test_tensor - &decompressed)?
            .pow(2.0)?
            .mean(&[], false)?
            .to_scalar::<f32>()?;

        // Measure compression ratio
        let original_size = test_tensor.numel() * 4; // f32
        let stats = storage.get_tensor_stats("test")?;
        let compression_ratio = original_size as f64 / stats.compressed_size as f64;

        results.push((
            name,
            compression_time.as_millis() as f64,
            decompression_time.as_millis() as f64,
            compression_ratio,
            mse,
        ));

        println!(
            "Level {}: compress={:.1}ms, decompress={:.1}ms, ratio={:.2}x, mse={:.6}",
            name,
            compression_time.as_millis() as f64,
            decompression_time.as_millis() as f64,
            compression_ratio,
            mse
        );

        // Verify time budget is respected (with 50% margin)
        assert!(compression_time.as_millis() as f64 < time_budget_ms * 1.5,
                "Compression time should respect budget");
    }

    // Verify tradeoff relationships
    // Higher compression levels should generally:
    // 1. Take more time
    // 2. Achieve better compression ratios
    // 3. Have better quality (lower MSE)

    for i in 1..results.len() {
        let (prev_name, prev_comp_time, _, prev_ratio, prev_mse) = results[i - 1];
        let (curr_name, curr_comp_time, _, curr_ratio, curr_mse) = results[i];

        // Allow some variance due to algorithm differences
        println!("Comparing {} vs {}", prev_name, curr_name);

        // Generally expect better compression ratios at higher levels
        if curr_ratio < prev_ratio * 0.8 {
            println!(
                "Warning: {} has lower compression ratio than {}",
                curr_name, prev_name
            );
        }
    }

    Ok(())
}

fn calculate_model_size(model: &HashMap<String, Tensor>) -> usize {
    model
        .values()
        .map(|t| t.numel() * t.dtype().size_in_bytes())
        .sum()
}