//! Tests for model pruning functionality
//!
//! Following strict TDD - tests define expected behavior before implementation
#![cfg(feature = "disabled_tests")]

use approx::assert_abs_diff_eq;
use rtx_compress::{
    Result,
    pruning::{
        MagnitudePruner, PruningConfig, PruningGranularity, PruningMethod, StructuredPruner,
    },
};
use rtx_tensor::{Device, Tensor};

#[test]
fn test_magnitude_pruner_creation() {
    let device = Device::cpu();

    let config = PruningConfig {
        sparsity: 0.5, // 50% sparsity
        structured: false,
        granularity: PruningGranularity::Unstructured,
        preserve_gradients: true,
    };

    let pruner = MagnitudePruner::new(config, &device);
    assert!(pruner.is_ok());

    let pruner = pruner.unwrap();
    assert_eq!(pruner.sparsity(), 0.5);
    assert!(!pruner.is_structured());
}

#[test]
fn test_unstructured_magnitude_pruning() {
    let device = Device::cpu();

    // Create a weight tensor with known values
    let weights_data = vec![
        0.1, -0.5, 0.3, -0.2, // Small and large magnitudes
        0.9, -0.05, 0.7, -0.8, // Mixed values
        0.01, -0.02, 0.6, -0.4, // Very small and medium
        0.15, -0.95, 0.03, -0.35,
    ];
    let weights = Tensor::from_data(weights_data.clone(), &[4, 4], &device).unwrap();

    // Prune 50% of weights (8 out of 16)
    let config = PruningConfig {
        sparsity: 0.5,
        structured: false,
        granularity: PruningGranularity::Unstructured,
        preserve_gradients: false,
    };

    let pruner = MagnitudePruner::new(config, &device).unwrap();
    let mask = pruner.compute_mask(&weights).unwrap();

    // Apply mask
    let pruned = pruner.apply_mask(&weights, &mask).unwrap();

    // Check that exactly 50% of weights are zero
    let pruned_data = pruned.to_vec().unwrap();
    let num_zeros = pruned_data.iter().filter(|&&x| x == 0.0).count();
    assert_eq!(num_zeros, 8);

    // Check that smallest magnitude weights were pruned
    // Values with smallest magnitude: 0.01, -0.02, 0.03, -0.05, 0.1, 0.15, -0.2, 0.3
    // Should keep: -0.5, 0.9, 0.7, -0.8, 0.6, -0.4, -0.95, -0.35
    for (i, val) in pruned_data.iter().enumerate() {
        let original = weights_data[i];
        if original.abs() < 0.35 {
            assert_eq!(
                *val, 0.0,
                "Small magnitude weight at index {} should be pruned",
                i
            );
        } else {
            assert_abs_diff_eq!(*val, original, epsilon = 1e-6);
        }
    }
}

#[test]
fn test_structured_pruning_channels() {
    let device = Device::cpu();

    // Create a Conv2d weight tensor [out_channels, in_channels, kernel_h, kernel_w]
    let weights = Tensor::randn(&[8, 3, 3, 3], &device).unwrap();

    // Prune 50% of output channels (4 out of 8)
    let config = PruningConfig {
        sparsity: 0.5,
        structured: true,
        granularity: PruningGranularity::Channel,
        preserve_gradients: false,
    };

    let pruner = StructuredPruner::new(config, &device).unwrap();
    let mask = pruner.compute_mask(&weights).unwrap();
    let pruned = pruner.apply_mask(&weights, &mask).unwrap();

    // Check that entire channels are pruned
    let pruned_shape = pruned.shape().dims();
    assert_eq!(pruned_shape, &[8, 3, 3, 3]);

    // Count pruned channels
    let pruned_data = pruned.to_vec().unwrap();
    let channel_size = 3 * 3 * 3;
    let mut pruned_channels = 0;

    for ch in 0..8 {
        let start = ch * channel_size;
        let end = start + channel_size;
        let channel_sum: f32 = pruned_data[start..end].iter().map(|x| x.abs()).sum();
        if channel_sum == 0.0 {
            pruned_channels += 1;
        }
    }

    assert_eq!(pruned_channels, 4, "Should prune exactly 4 channels");
}

#[test]
fn test_structured_pruning_filters() {
    let device = Device::cpu();

    // Create weight tensor for structured filter pruning
    let weights = Tensor::randn(&[64, 64, 3, 3], &device).unwrap();

    let config = PruningConfig {
        sparsity: 0.25, // Prune 25% of filters
        structured: true,
        granularity: PruningGranularity::Filter,
        preserve_gradients: false,
    };

    let pruner = StructuredPruner::new(config, &device).unwrap();
    let indices = pruner.get_pruned_indices(&weights).unwrap();

    // Should identify 16 filters to prune (25% of 64)
    assert_eq!(indices.len(), 16);

    // Apply pruning
    let pruned = pruner.prune(&weights).unwrap();

    // Verify filters are zeroed
    let pruned_data = pruned.to_vec().unwrap();
    let filter_size = 64 * 3 * 3;

    for &idx in indices.iter() {
        let start = idx * filter_size;
        let end = start + filter_size;
        let filter_sum: f32 = pruned_data[start..end].iter().map(|x| x.abs()).sum();
        assert_eq!(
            filter_sum, 0.0,
            "Filter {} should be completely pruned",
            idx
        );
    }
}

#[test]
fn test_gradual_magnitude_pruning() {
    let device = Device::cpu();

    // Test gradual pruning over multiple iterations
    let weights = Tensor::randn(&[100, 100], &device).unwrap();

    let mut pruner = MagnitudePruner::with_schedule(
        0.9, // Final sparsity
        100, // Number of steps
        &device,
    )
    .unwrap();

    // Initial pruning should be minimal
    let mask_initial = pruner.compute_mask_at_step(&weights, 0).unwrap();
    let sparsity_initial = pruner.get_current_sparsity(0);
    assert!(sparsity_initial < 0.1, "Initial sparsity should be low");

    // Middle of training
    let mask_middle = pruner.compute_mask_at_step(&weights, 50).unwrap();
    let sparsity_middle = pruner.get_current_sparsity(50);
    assert!(
        sparsity_middle > 0.3 && sparsity_middle < 0.6,
        "Middle sparsity should be moderate"
    );

    // End of training
    let mask_final = pruner.compute_mask_at_step(&weights, 99).unwrap();
    let sparsity_final = pruner.get_current_sparsity(99);
    assert_abs_diff_eq!(sparsity_final, 0.9, epsilon = 0.01);
}

#[test]
fn test_pruning_with_importance_scores() {
    let device = Device::cpu();

    // Create weights and importance scores
    let weights = Tensor::randn(&[10, 10], &device).unwrap();
    let importance = Tensor::rand(&[10, 10], &device).unwrap(); // Random importance scores

    let config = PruningConfig {
        sparsity: 0.3,
        structured: false,
        granularity: PruningGranularity::Unstructured,
        preserve_gradients: false,
    };

    let pruner = MagnitudePruner::new(config, &device).unwrap();

    // Prune based on importance scores instead of magnitude
    let mask = pruner
        .compute_mask_with_importance(&weights, &importance)
        .unwrap();
    let pruned = pruner.apply_mask(&weights, &mask).unwrap();

    // Check sparsity
    let pruned_data = pruned.to_vec().unwrap();
    let num_zeros = pruned_data.iter().filter(|&&x| x == 0.0).count();
    let expected_zeros = (100.0 * 0.3) as usize;
    assert_eq!(num_zeros, expected_zeros);
}

#[test]
fn test_pruning_mask_persistence() {
    let device = Device::cpu();

    let config = PruningConfig {
        sparsity: 0.4,
        structured: false,
        granularity: PruningGranularity::Unstructured,
        preserve_gradients: true,
    };

    let pruner = MagnitudePruner::new(config, &device).unwrap();
    let weights = Tensor::randn(&[50, 50], &device).unwrap();

    // Compute and save mask
    let mask = pruner.compute_mask(&weights).unwrap();
    let mask_data = mask.to_vec().unwrap();

    // Apply mask multiple times - should be consistent
    let pruned1 = pruner.apply_mask(&weights, &mask).unwrap();
    let pruned2 = pruner.apply_mask(&weights, &mask).unwrap();

    let data1 = pruned1.to_vec().unwrap();
    let data2 = pruned2.to_vec().unwrap();

    for i in 0..data1.len() {
        assert_abs_diff_eq!(data1[i], data2[i], epsilon = 1e-6);
    }
}

#[test]
fn test_fine_grained_pruning() {
    let device = Device::cpu();

    // Test N:M sparsity pattern (e.g., 2:4 for NVIDIA Ampere)
    let weights = Tensor::randn(&[64, 64], &device).unwrap();

    let pruner = MagnitudePruner::with_nm_sparsity(
        2, // N (number of non-zeros)
        4, // M (block size)
        &device,
    )
    .unwrap();

    let mask = pruner.compute_mask(&weights).unwrap();
    let pruned = pruner.apply_mask(&weights, &mask).unwrap();

    // Verify 2:4 pattern
    let pruned_data = pruned.to_vec().unwrap();

    // Check each block of 4 elements has exactly 2 zeros
    for i in (0..pruned_data.len()).step_by(4) {
        if i + 3 < pruned_data.len() {
            let block = &pruned_data[i..i + 4];
            let zeros = block.iter().filter(|&&x| x == 0.0).count();
            assert_eq!(zeros, 2, "Each block of 4 should have exactly 2 zeros");
        }
    }
}

#[test]
fn test_pruning_statistics() {
    let device = Device::cpu();

    let config = PruningConfig {
        sparsity: 0.6,
        structured: false,
        granularity: PruningGranularity::Unstructured,
        preserve_gradients: false,
    };

    let pruner = MagnitudePruner::new(config, &device).unwrap();
    let weights = Tensor::randn(&[128, 128], &device).unwrap();

    let stats = pruner.analyze(&weights).unwrap();

    // Check statistics
    assert_eq!(stats.total_parameters, 128 * 128);
    assert_eq!(stats.target_sparsity, 0.6);
    assert!(stats.achieved_sparsity >= 0.0 && stats.achieved_sparsity <= 1.0);
    assert!(stats.compression_ratio > 1.0); // Should achieve some compression
}