use rtx_compress::{
    CompressionError, Result,
    pipeline::{
        CompressionPipeline, CompressionPipelineConfig, CompressionStrategy, HardwareTarget,
    },
    pruning::{
        ImportanceMetric, PruningCriterion, SparsityPattern, StructuredPruner,
        StructuredPruningConfig, StructuredPruningMethod, UnstructuredPruner,
        UnstructuredPruningConfig,
    },
    quantization::{PQConfig, ProductQuantizer, VQConfig, VectorQuantizer},
};
use rtx_tensor::{Device, Tensor};
use std::collections::HashMap;

// Import enums for VectorQuantizer
use rtx_compress::quantization::vector_quantization::{CodebookInitialization, DistanceMetric};

#[test]
#[ignore = "Pre-existing quantization codebook size issue"]
fn test_product_quantization_basic() {
    let device = Device::cpu();
    let config = PQConfig {
        num_subquantizers: 4,
        codebook_size: 256,
        max_iterations: 10,
        tolerance: 1e-6,
        use_opq: false,
        opq_iterations: 10,
        use_residual: false,
        residual_stages: 0,
    };

    let mut pq = ProductQuantizer::new(config).unwrap();

    // Create test data
    let data = Tensor::randn(&[100, 16], &device).unwrap();

    // Fit the quantizer
    pq.fit(&data).unwrap();

    // Encode and decode
    let codes = pq.encode(&data).unwrap();
    let reconstructed = pq.decode(&codes).unwrap();

    // Check shapes
    assert_eq!(reconstructed.shape().dims(), data.shape().dims());
}

#[test]
#[ignore = "Pre-existing quantization codebook size issue"]
fn test_vector_quantization() {
    let device = Device::cpu();
    let config = VQConfig {
        codebook_size: 512,
        vector_dim: 64,
        max_iterations: 100,
        tolerance: 1e-6,
        initialization:
            rtx_compress::quantization::vector_quantization::CodebookInitialization::Random,
        distance_metric: rtx_compress::quantization::vector_quantization::DistanceMetric::Euclidean,
    };

    let mut vq = VectorQuantizer::new(config);

    // Create test data
    let data = Tensor::randn(&[32, 64], &device).unwrap();

    // Train the quantizer
    vq.fit(&data).unwrap();

    // Quantize data
    let codes = vq.encode(&data).unwrap();

    assert_eq!(codes.shape().dims()[0], data.shape().dims()[0]);
}

#[test]
fn test_compression_pipeline() {
    let device = Device::cpu();
    let config = CompressionPipelineConfig {
        strategy: CompressionStrategy::Balanced,
        target_compression_ratio: 2.0,
        target_accuracy_retention: 0.95,
        progressive: false,
        num_stages: 1,
        validation_size: 100,
        hardware_target: HardwareTarget::CPU,
    };

    let mut pipeline = CompressionPipeline::new(config).unwrap();

    // Create test model parameters
    let mut params = HashMap::new();
    params.insert(
        "weight".to_string(),
        Tensor::randn(&[512, 512], &device).unwrap(),
    );
    params.insert("bias".to_string(), Tensor::randn(&[512], &device).unwrap());

    // Compress the model
    let result = pipeline.compress(&params, None, None).unwrap();
    let compressed = result.compressed_model;

    // Verify keys match
    assert_eq!(params.len(), compressed.len());
    for key in params.keys() {
        assert!(compressed.contains_key(key));
    }
}

#[test]
fn test_structured_pruning() {
    let device = Device::cpu();
    let config = StructuredPruningConfig {
        method: StructuredPruningMethod::ChannelPruning,
        importance_metric: ImportanceMetric::L2Norm,
        target_sparsity: 0.5,
        schedule: None,
        block_size: 4,
        min_channels: None,
        hardware_alignment: None,
        gpu_optimized: false,
        layer_wise_ratios: None,
        generate_masks: false,
        use_distillation: false,
        distillation_weight: 0.0,
        recovery_epochs: 0,
        recovery_learning_rate: 0.001,
    };
    let pruner = StructuredPruner::new(config).unwrap();

    // Create test weight tensor
    let weight = Tensor::randn(&[64, 128], &device).unwrap();

    // Apply pruning with 50% sparsity
    let result = pruner.prune_tensor(&weight, "test.weight").unwrap();

    // Check shape is preserved for channel pruning (reduces output channels)
    assert!(result.shape().dims()[0] <= weight.shape().dims()[0]);
}

#[test]
fn test_unstructured_pruning() {
    let device = Device::cpu();
    let config = UnstructuredPruningConfig {
        criterion: PruningCriterion::GlobalMagnitude,
        target_sparsity: 0.9,
        layer_wise_ratios: None,
        schedule: None,
        pattern: SparsityPattern::Unstructured,
        layer_sensitivities: None,
        generate_masks: false,
        collect_statistics: false,
        min_threshold: 0.0,
    };
    let pruner = UnstructuredPruner::new(config).unwrap();

    // Create test tensor
    let tensor = Tensor::randn(&[256, 256], &device).unwrap();

    // Apply magnitude-based pruning
    let mut model = HashMap::new();
    model.insert("weight".to_string(), tensor.clone());

    let result = pruner.prune_model(&model).unwrap();
    let pruned = &result["weight"];

    assert_eq!(pruned.shape().dims(), tensor.shape().dims());
}

#[test]
fn test_quantization_error_handling() {
    let device = Device::cpu();
    let config = PQConfig {
        num_subquantizers: 4,
        codebook_size: 256,
        max_iterations: 10,
        tolerance: 1e-6,
        use_opq: false,
        opq_iterations: 10,
        use_residual: false,
        residual_stages: 0,
    };

    let pq = ProductQuantizer::new(config).unwrap();

    // Try to encode without training - should fail
    let data = Tensor::randn(&[10, 16], &device).unwrap();
    let result = pq.encode(&data);

    assert!(result.is_err());
}

#[test]
fn test_compression_with_calibration_data() {
    let device = Device::cpu();
    let config = CompressionPipelineConfig {
        strategy: CompressionStrategy::Accuracy,
        target_compression_ratio: 2.0,
        target_accuracy_retention: 0.98,
        progressive: false,
        num_stages: 1,
        validation_size: 100,
        hardware_target: HardwareTarget::CPU,
    };

    let mut pipeline = CompressionPipeline::new(config).unwrap();

    // Now compress with calibrated settings
    let params = HashMap::from([(
        "layer1".to_string(),
        Tensor::randn(&[512, 512], &device).unwrap(),
    )]);

    let result = pipeline.compress(&params, None, None).unwrap();
    assert!(!result.compressed_model.is_empty());
}

#[test]
fn test_mixed_precision_optimization() {
    let device = Device::cpu();
    let config = CompressionPipelineConfig {
        strategy: CompressionStrategy::Speed,
        target_compression_ratio: 1.5,
        target_accuracy_retention: 0.99,
        progressive: false,
        num_stages: 1,
        validation_size: 100,
        hardware_target: HardwareTarget::GPU,
    };

    let pipeline = CompressionPipeline::new(config).unwrap();

    // Test pipeline is properly configured - just verify it was created
    assert!(true);
}

#[test]
#[ignore = "Pre-existing quantization codebook size issue"]
fn test_memory_efficiency() {
    let device = Device::cpu();
    let config = PQConfig {
        num_subquantizers: 8,
        codebook_size: 256,
        max_iterations: 5,
        tolerance: 1e-6,
        use_opq: true,
        opq_iterations: 10,
        use_residual: false,
        residual_stages: 0,
    };

    let mut pq = ProductQuantizer::new(config).unwrap();

    // Large dataset
    let data = Tensor::randn(&[1000, 128], &device).unwrap();

    // Should handle large data efficiently
    pq.fit(&data).unwrap();
    let codes = pq.encode(&data).unwrap();

    // Codes should be much smaller than original data
    let code_size = codes.shape().dims().iter().product::<usize>();
    let data_size = data.shape().dims().iter().product::<usize>();

    assert!(code_size < data_size);
}