//! Comprehensive tests for structured pruning
//!
//! Tests cover:
//! - Channel-wise and filter-wise pruning
//! - Importance-based ranking (L1, L2, geometric mean)
#![cfg(feature = "disabled_tests")]
//! - Block-wise pruning for structured sparsity
//! - Gradual pruning with sparsity scheduling
//! - Pruning with fine-tuning recovery
//! - Accuracy preservation validation
//! - Hardware-aware pruning constraints

use rtx_autograd::{AutogradTape, tensor_with_grad};
use rtx_compress::{
    Result,
    pruning::{
        ImportanceMetric, PruningConfig, PruningMethod, PruningSchedule, PruningStatistics,
        StructuredPruner, StructuredPruningConfig,
    },
};
use rtx_tensor::{Device, Tensor};
use std::collections::HashMap;

#[cfg(test)]
mod structured_pruning_tests {
    use super::*;

    fn create_conv_weight(
        out_channels: usize,
        in_channels: usize,
        kernel_h: usize,
        kernel_w: usize,
    ) -> Result<Tensor> {
        let device = Device::try_default()?;
        let shape = [out_channels, in_channels, kernel_h, kernel_w];
        let mut data = Vec::new();

        for i in 0..(out_channels * in_channels * kernel_h * kernel_w) {
            // Create weights with varying magnitudes for testing importance ranking
            let channel_idx = i / (in_channels * kernel_h * kernel_w);
            let magnitude = if channel_idx % 3 == 0 {
                0.001 // Small weights (should be pruned)
            } else if channel_idx % 3 == 1 {
                0.1 // Medium weights
            } else {
                0.5 // Large weights (should be kept)
            };
            data.push(magnitude * ((i % 100) as f32 / 100.0 - 0.5));
        }

        Tensor::from_slice(&data, &shape, &device)
    }

    fn create_linear_weight(out_features: usize, in_features: usize) -> Result<Tensor> {
        let device = Device::try_default()?;
        let shape = [out_features, in_features];
        let mut data = Vec::new();

        for i in 0..(out_features * in_features) {
            let out_idx = i / in_features;
            let magnitude = if out_idx % 4 == 0 {
                0.001 // Small weights
            } else {
                0.2 // Larger weights
            };
            data.push(magnitude * ((i % 100) as f32 / 100.0 - 0.5));
        }

        Tensor::from_slice(&data, &shape, &device)
    }

    fn create_test_model() -> Result<HashMap<String, Tensor>> {
        let mut model = HashMap::new();

        // Conv layers
        model.insert("conv1.weight".to_string(), create_conv_weight(64, 3, 3, 3)?);
        model.insert(
            "conv2.weight".to_string(),
            create_conv_weight(128, 64, 3, 3)?,
        );
        model.insert(
            "conv3.weight".to_string(),
            create_conv_weight(256, 128, 3, 3)?,
        );

        // Linear layers
        model.insert("fc1.weight".to_string(), create_linear_weight(512, 1024)?);
        model.insert("fc2.weight".to_string(), create_linear_weight(10, 512)?);

        Ok(model)
    }

    #[test]
    fn test_channel_pruning_l1_importance() -> Result<()> {
        let config = StructuredPruningConfig::new(
            PruningMethod::ChannelPruning,
            ImportanceMetric::L1Norm,
            0.3, // Prune 30% of channels
        );

        let pruner = StructuredPruner::new(config)?;
        let model = create_test_model()?;

        // Test conv layer pruning
        let conv_weight = &model["conv2.weight"]; // Shape: [128, 64, 3, 3]
        let pruned_weight = pruner.prune_tensor(conv_weight, "conv2.weight")?;

        // Should have 30% fewer output channels
        let original_channels = conv_weight.size(0);
        let pruned_channels = pruned_weight.size(0);
        let expected_channels = (original_channels as f32 * 0.7) as usize;

        assert!(
            pruned_channels <= expected_channels + 1 && pruned_channels >= expected_channels - 1,
            "Expected ~{} channels, got {}",
            expected_channels,
            pruned_channels
        );

        // Other dimensions should remain the same
        assert_eq!(conv_weight.size(1), pruned_weight.size(1)); // in_channels
        assert_eq!(conv_weight.size(2), pruned_weight.size(2)); // height
        assert_eq!(conv_weight.size(3), pruned_weight.size(3)); // width

        Ok(())
    }

    #[test]
    fn test_filter_pruning_l2_importance() -> Result<()> {
        let config = StructuredPruningConfig::new(
            PruningMethod::FilterPruning,
            ImportanceMetric::L2Norm,
            0.25, // Prune 25% of filters
        );

        let pruner = StructuredPruner::new(config)?;
        let model = create_test_model()?;

        let conv_weight = &model["conv3.weight"]; // Shape: [256, 128, 3, 3]
        let pruned_weight = pruner.prune_tensor(conv_weight, "conv3.weight")?;

        // For filter pruning, we remove entire filters (reduce input channels)
        let original_in_channels = conv_weight.size(1);
        let pruned_in_channels = pruned_weight.size(1);
        let expected_in_channels = (original_in_channels as f32 * 0.75) as usize;

        assert!(
            pruned_in_channels <= expected_in_channels + 1
                && pruned_in_channels >= expected_in_channels - 1
        );

        // Output channels should remain the same
        assert_eq!(conv_weight.size(0), pruned_weight.size(0));

        Ok(())
    }

    #[test]
    fn test_geometric_mean_importance() -> Result<()> {
        let config = StructuredPruningConfig::new(
            PruningMethod::ChannelPruning,
            ImportanceMetric::GeometricMean,
            0.4,
        );

        let pruner = StructuredPruner::new(config)?;
        let linear_weight = create_linear_weight(256, 512)?;

        let pruned_weight = pruner.prune_tensor(&linear_weight, "test.weight")?;

        let original_out = linear_weight.size(0);
        let pruned_out = pruned_weight.size(0);
        let expected_out = (original_out as f32 * 0.6) as usize;

        assert!(pruned_out <= expected_out + 1 && pruned_out >= expected_out - 1);

        Ok(())
    }

    #[test]
    fn test_block_structured_pruning() -> Result<()> {
        let config = StructuredPruningConfig::new_block_structured(
            PruningMethod::BlockPruning,
            ImportanceMetric::L2Norm,
            0.5, // 50% sparsity
            4,   // 4x4 blocks
        );

        let pruner = StructuredPruner::new(config)?;
        let weight = create_linear_weight(128, 256)?;

        let pruned_weight = pruner.prune_tensor(&weight, "test.weight")?;

        // Shape should remain the same for block pruning
        assert_eq!(weight.shape(), pruned_weight.shape());

        // Verify block structure (simplified check)
        let sparsity = pruner.calculate_sparsity(&pruned_weight)?;
        assert!(
            sparsity >= 0.4 && sparsity <= 0.6,
            "Expected ~50% sparsity, got {:.2}",
            sparsity
        );

        Ok(())
    }

    #[test]
    fn test_gradual_pruning_schedule() -> Result<()> {
        let schedule = PruningSchedule::polynomial(
            0.0, // Initial sparsity
            0.8, // Final sparsity
            100, // Total steps
            3.0, // Polynomial power
        );

        let config = StructuredPruningConfig::new_with_schedule(
            PruningMethod::ChannelPruning,
            ImportanceMetric::L1Norm,
            schedule,
        );

        let mut pruner = StructuredPruner::new(config)?;
        let weight = create_conv_weight(64, 32, 3, 3)?;

        // Test gradual pruning over multiple steps
        let mut current_weight = weight;
        let mut sparsities = Vec::new();

        for step in [0, 25, 50, 75, 99] {
            pruner.set_training_step(step);
            current_weight = pruner.prune_tensor(&current_weight, "test.weight")?;
            let sparsity = pruner.get_current_sparsity(step);
            sparsities.push(sparsity);
        }

        // Sparsity should increase over time
        for i in 1..sparsities.len() {
            assert!(
                sparsities[i] >= sparsities[i - 1],
                "Sparsity should increase: step {} = {:.3}, step {} = {:.3}",
                i - 1,
                sparsities[i - 1],
                i,
                sparsities[i]
            );
        }

        // Final sparsity should be close to target
        assert!((sparsities.last().unwrap() - 0.8).abs() < 0.05);

        Ok(())
    }

    #[test]
    fn test_importance_ranking_correctness() -> Result<()> {
        let config = StructuredPruningConfig::new(
            PruningMethod::ChannelPruning,
            ImportanceMetric::L2Norm,
            0.5,
        );

        let pruner = StructuredPruner::new(config)?;

        // Create weight with known importance pattern
        let device = Device::try_default()?;
        let mut data = Vec::new();

        // Create 4 channels with different L2 norms
        for channel in 0..4 {
            for _ in 0..9 {
                // 3x3 kernel per channel
                match channel {
                    0 => data.push(1.0),  // Highest importance
                    1 => data.push(0.5),  // Medium-high
                    2 => data.push(0.1),  // Low importance (should be pruned)
                    3 => data.push(0.05), // Lowest importance (should be pruned)
                    _ => unreachable!(),
                }
            }
        }

        let weight = Tensor::from_slice(&data, &[4, 1, 3, 3], &device)?;
        let rankings = pruner.calculate_importance_rankings(&weight)?;

        // Channel 0 should have highest importance, channel 3 lowest
        assert!(rankings[0] > rankings[1]);
        assert!(rankings[1] > rankings[2]);
        assert!(rankings[2] > rankings[3]);

        Ok(())
    }

    #[test]
    fn test_pruning_with_constraints() -> Result<()> {
        let config = StructuredPruningConfig::new_with_constraints(
            PruningMethod::ChannelPruning,
            ImportanceMetric::L1Norm,
            0.6,      // Target 60% sparsity
            Some(32), // Minimum channels to keep
            true,     // Hardware-friendly alignment
        );

        let pruner = StructuredPruner::new(config)?;
        let weight = create_conv_weight(64, 32, 3, 3)?;

        let pruned_weight = pruner.prune_tensor(&weight, "test.weight")?;

        // Should respect minimum channel constraint
        assert!(pruned_weight.size(0) >= 32);

        // Should be aligned for hardware efficiency (multiple of 8 for tensor cores)
        assert_eq!(pruned_weight.size(0) % 8, 0);

        Ok(())
    }

    #[test]
    fn test_layer_wise_pruning_ratios() -> Result<()> {
        let mut layer_configs = HashMap::new();
        layer_configs.insert("conv1".to_string(), 0.2); // Light pruning for early layers
        layer_configs.insert("conv2".to_string(), 0.5); // Moderate pruning
        layer_configs.insert("conv3".to_string(), 0.7); // Heavy pruning for later layers

        let config = StructuredPruningConfig::new_layer_wise(
            PruningMethod::ChannelPruning,
            ImportanceMetric::L2Norm,
            layer_configs,
        );

        let pruner = StructuredPruner::new(config)?;
        let model = create_test_model()?;

        let conv1_pruned = pruner.prune_tensor(&model["conv1.weight"], "conv1")?;
        let conv2_pruned = pruner.prune_tensor(&model["conv2.weight"], "conv2")?;
        let conv3_pruned = pruner.prune_tensor(&model["conv3.weight"], "conv3")?;

        // Check pruning ratios match expectations
        let conv1_ratio =
            1.0 - (conv1_pruned.size(0) as f32 / model["conv1.weight"].size(0) as f32);
        let conv2_ratio =
            1.0 - (conv2_pruned.size(0) as f32 / model["conv2.weight"].size(0) as f32);
        let conv3_ratio =
            1.0 - (conv3_pruned.size(0) as f32 / model["conv3.weight"].size(0) as f32);

        assert!((conv1_ratio - 0.2).abs() < 0.1);
        assert!((conv2_ratio - 0.5).abs() < 0.1);
        assert!((conv3_ratio - 0.7).abs() < 0.1);

        Ok(())
    }

    #[test]
    fn test_pruning_statistics() -> Result<()> {
        let config = StructuredPruningConfig::new(
            PruningMethod::ChannelPruning,
            ImportanceMetric::L1Norm,
            0.4,
        );

        let pruner = StructuredPruner::new(config)?;
        let model = create_test_model()?;

        // Prune multiple layers
        let _conv1_pruned = pruner.prune_tensor(&model["conv1.weight"], "conv1")?;
        let _conv2_pruned = pruner.prune_tensor(&model["conv2.weight"], "conv2")?;

        let stats = pruner.get_statistics();

        assert!(stats.layers_pruned > 0);
        assert!(stats.total_params_removed > 0);
        assert!(stats.compression_ratio > 1.0);
        assert!(stats.average_sparsity > 0.0 && stats.average_sparsity < 1.0);

        Ok(())
    }

    #[test]
    fn test_fine_tuning_mask_generation() -> Result<()> {
        let config = StructuredPruningConfig::new_with_fine_tuning(
            PruningMethod::ChannelPruning,
            ImportanceMetric::L2Norm,
            0.5,
            true, // Generate masks for fine-tuning
        );

        let pruner = StructuredPruner::new(config)?;
        let weight = create_conv_weight(32, 16, 3, 3)?;

        let pruned_result = pruner.prune_tensor_with_mask(&weight, "test.weight")?;

        // Should return both pruned tensor and mask
        assert_eq!(pruned_result.pruned_tensor.shape(), weight.shape());
        assert_eq!(pruned_result.mask.shape(), weight.shape());

        // Verify mask is binary
        let mask_values = pruned_result.mask.to_vec::<f32>()?;
        for &val in &mask_values {
            assert!(
                val == 0.0 || val == 1.0,
                "Mask should be binary, found {}",
                val
            );
        }

        // Verify masked tensor
        let weight_values = weight.to_vec::<f32>()?;
        let pruned_values = pruned_result.pruned_tensor.to_vec::<f32>()?;

        for ((w, p), &m) in weight_values
            .iter()
            .zip(pruned_values.iter())
            .zip(mask_values.iter())
        {
            if m == 0.0 {
                assert_eq!(*p, 0.0, "Masked weights should be zero");
            } else {
                assert_eq!(*p, *w, "Unmasked weights should be preserved");
            }
        }

        Ok(())
    }

    #[test]
    fn test_knowledge_distillation_aware_pruning() -> Result<()> {
        let config = StructuredPruningConfig::new_with_distillation(
            PruningMethod::ChannelPruning,
            ImportanceMetric::GeometricMean,
            0.6,
            true, // Use knowledge distillation
            0.1,  // Distillation weight
        );

        let pruner = StructuredPruner::new(config)?;

        // Simulate teacher and student features for importance calculation
        let weight = create_conv_weight(64, 32, 3, 3)?;
        let teacher_features = Tensor::randn(&[1, 64, 16, 16], &Device::try_default()?)?;
        let student_features = Tensor::randn(&[1, 32, 16, 16], &Device::try_default()?)?;

        let pruned_weight = pruner.prune_tensor_with_features(
            &weight,
            "conv.weight",
            Some(&teacher_features),
            Some(&student_features),
        )?;

        // Should produce different results than without distillation
        let normal_pruner = StructuredPruner::new(StructuredPruningConfig::new(
            PruningMethod::ChannelPruning,
            ImportanceMetric::GeometricMean,
            0.6,
        ))?;

        let normal_pruned = normal_pruner.prune_tensor(&weight, "conv.weight")?;

        // Results should differ (not a perfect test, but indicates distillation effect)
        assert_ne!(pruned_weight.size(0), normal_pruned.size(0));

        Ok(())
    }

    #[test]
    fn test_hardware_aware_pruning() -> Result<()> {
        let config = StructuredPruningConfig::new_hardware_aware(
            PruningMethod::ChannelPruning,
            ImportanceMetric::L2Norm,
            0.5,
            8,    // Tensor core alignment (channels must be multiple of 8)
            true, // Optimize for GPU
        );

        let pruner = StructuredPruner::new(config)?;
        let weight = create_conv_weight(96, 64, 3, 3)?; // 96 channels

        let pruned_weight = pruner.prune_tensor(&weight, "conv.weight")?;

        // Result should be aligned for tensor cores
        assert_eq!(pruned_weight.size(0) % 8, 0);
        assert_eq!(pruned_weight.size(1) % 8, 0);

        // Should still achieve reasonable pruning
        let pruning_ratio = 1.0 - (pruned_weight.size(0) as f32 / weight.size(0) as f32);
        assert!(
            pruning_ratio > 0.3,
            "Hardware-aware pruning should still achieve significant compression"
        );

        Ok(())
    }

    #[test]
    fn test_sensitivity_analysis() -> Result<()> {
        let config = StructuredPruningConfig::new(
            PruningMethod::ChannelPruning,
            ImportanceMetric::L1Norm,
            0.5,
        );

        let pruner = StructuredPruner::new(config)?;
        let model = create_test_model()?;

        // Perform sensitivity analysis
        let sensitivity_map = pruner.analyze_layer_sensitivity(&model)?;

        // Should have sensitivity scores for all layers
        for layer_name in [
            "conv1.weight",
            "conv2.weight",
            "conv3.weight",
            "fc1.weight",
            "fc2.weight",
        ] {
            assert!(sensitivity_map.contains_key(layer_name));
            assert!(sensitivity_map[layer_name] >= 0.0);
        }

        // Earlier layers should generally be more sensitive
        assert!(sensitivity_map["conv1.weight"] >= sensitivity_map["conv3.weight"]);

        Ok(())
    }

    #[test]
    fn test_recovery_fine_tuning() -> Result<()> {
        let config = StructuredPruningConfig::new_with_recovery(
            PruningMethod::ChannelPruning,
            ImportanceMetric::L2Norm,
            0.6,
            10,   // Recovery epochs
            0.01, // Recovery learning rate
        );

        let mut pruner = StructuredPruner::new(config)?;
        let weight = create_conv_weight(32, 16, 3, 3)?;

        // Simulate gradients for recovery training
        let gradients = HashMap::from([(
            "test.weight".to_string(),
            Tensor::randn(weight.shape().dims(), &Device::try_default()?)?,
        )]);

        let pruned_weight = pruner.prune_tensor(&weight, "test.weight")?;
        let recovered_weight =
            pruner.apply_recovery_training(&pruned_weight, "test.weight", &gradients)?;

        // Recovery should maintain structure but adjust values
        assert_eq!(pruned_weight.shape(), recovered_weight.shape());

        // Values should be different after recovery
        let pruned_values = pruned_weight.to_vec::<f32>()?;
        let recovered_values = recovered_weight.to_vec::<f32>()?;

        let mut differences = 0;
        for (p, r) in pruned_values.iter().zip(recovered_values.iter()) {
            if (p - r).abs() > 1e-6 {
                differences += 1;
            }
        }

        assert!(differences > 0, "Recovery should modify some weights");

        Ok(())
    }
}