// Basic compression tests - simplified to match current API
use rtx_compress::{
    Result,
    pipeline::{
        CompressionPipeline, CompressionPipelineConfig, CompressionStrategy, HardwareTarget,
    },
    quantization::{
        CalibrationConfig, PostTrainingQuantizer, QuantizationConfig, QuantizationScheme,
    },
};
use rtx_tensor::{Device, Tensor};
use std::collections::HashMap;

#[test]
fn test_compression_pipeline_creation() {
    let config = CompressionPipelineConfig {
        strategy: CompressionStrategy::Balanced,
        target_compression_ratio: 4.0,
        target_accuracy_retention: 0.95,
        progressive: false,
        num_stages: 2,
        validation_size: 100,
        hardware_target: HardwareTarget::CPU,
    };
    let pipeline = CompressionPipeline::new(config);
    assert!(pipeline.is_ok(), "Should create compression pipeline");
}

#[test]
#[ignore = "Pre-existing Metal device randn issue"]
fn test_post_training_quantization() {
    let device = Device::try_default().unwrap();

    let calibration = CalibrationConfig::new(100, 0.99);
    let config = QuantizationConfig::new(QuantizationScheme::INT8, calibration);

    let quantizer = PostTrainingQuantizer::new(config);
    assert!(quantizer.is_ok(), "Should create post-training quantizer");

    let mut quantizer = quantizer.unwrap();

    // Create calibration data
    let calibration_data: Vec<Tensor> = (0..10)
        .map(|_| Tensor::randn(&[64, 64], &device).unwrap())
        .collect();

    let calibrate_result = quantizer.calibrate(&calibration_data);
    assert!(calibrate_result.is_ok(), "Should calibrate quantizer");

    // Quantize a tensor
    let tensor = Tensor::randn(&[64, 64], &device).unwrap();
    let quantized = quantizer.quantize(&tensor);
    assert!(quantized.is_ok(), "Should quantize tensor");

    let quantized_tensor = quantized.unwrap();

    // Dequantize
    let dequantized = quantizer.dequantize(&quantized_tensor);
    assert!(dequantized.is_ok(), "Should dequantize tensor");

    // Check shape preservation
    let deq = dequantized.unwrap();
    assert_eq!(deq.shape().dims(), tensor.shape().dims());
}

#[test]
fn test_compression_pipeline_auto() {
    let pipeline = CompressionPipeline::auto(4.0, 0.95, HardwareTarget::GPU);
    assert!(pipeline.is_ok(), "Should create auto-tuned pipeline");
}

#[test]
#[ignore = "Pre-existing Metal device randn issue"]
fn test_end_to_end_compression_pipeline() -> Result<()> {
    let device = Device::try_default()?;

    // Create a simple model
    let mut model = HashMap::new();
    model.insert("layer1".to_string(), Tensor::randn(&[512, 512], &device)?);
    model.insert("layer2".to_string(), Tensor::randn(&[512, 256], &device)?);
    model.insert("layer3".to_string(), Tensor::randn(&[256, 128], &device)?);

    let config = CompressionPipelineConfig {
        strategy: CompressionStrategy::Size,
        target_compression_ratio: 4.0,
        target_accuracy_retention: 0.90,
        progressive: false,
        num_stages: 1,
        validation_size: 100,
        hardware_target: HardwareTarget::CPU,
    };

    let mut pipeline = CompressionPipeline::new(config)?;

    // Apply compression
    let result = pipeline.compress(&model, None, None)?;

    // Verify compression was achieved
    assert!(
        result.statistics.compression_ratio >= 1.0,
        "Should achieve some compression"
    );

    Ok(())
}