//! Magnitude-based weight pruning for neural networks
//!
//! This module implements magnitude-based pruning techniques that remove parameters
//! with the smallest absolute values, a fundamental approach in neural network compression.
//! Magnitude pruning is based on the hypothesis that small weights contribute less to
//! the network's output and can be safely removed.
//!
//! ## Supported Pruning Types
//!
//! - **Global Magnitude Pruning**: Remove smallest weights across the entire network
//! - **Layer-wise Magnitude Pruning**: Remove smallest weights within each layer independently  
//! - **Gradual Pruning**: Apply pruning schedules during training (linear, polynomial, exponential)
//! - **Fine-grained Pruning**: Element-wise weight removal (unstructured)
//! - **Structured Magnitude Pruning**: Channel-wise, filter-wise pruning patterns
//!
//! ## Key Features
//!
//! - Configurable sparsity ratios (0-100%)
//! - Multiple pruning granularities (element, channel, filter)
//! - Pruning schedules for gradual sparsification during training
//! - Mask generation and application for efficient training
//! - Integration with existing sparse tensor infrastructure
//! - Support for both weight and activation pruning
//!
//! ## Design Principles
//!
//! - **Memory Efficiency**: Leverage sparse tensor formats for storage
//! - **Training Integration**: Seamless integration with optimizers and training loops
//! - **Hardware Acceleration**: Support for structured patterns that map to modern accelerators
//! - **Flexibility**: Support various pruning strategies and schedules
//!
//! ## Usage Example
//!
//! ```rust
//! use rtx_transformers::layers::{MagnitudePruning, PruningConfig, PruningStrategy};
//! use rtx_tensor::{Tensor, Device};
//!
//! let device = Device::cpu();
//! let weights = Tensor::randn(&[512, 1024], &device).unwrap();
//!
//! // Global magnitude pruning with 50% sparsity
//! let config = PruningConfig::global(0.5);
//! let pruned_weights = MagnitudePruning::prune(&weights, &config).unwrap();
//!
//! // Layer-wise pruning with gradual schedule
//! let schedule_config = PruningConfig::layerwise(0.8)
//!     .with_schedule(PruningSchedule::polynomial(100, 3.0));
//! let mask = MagnitudePruning::create_mask(&weights, &schedule_config).unwrap();
//! ```

use crate::{Result, TransformerError};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Configuration for magnitude-based pruning
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PruningConfig {
    /// Pruning strategy to apply
    pub strategy: PruningStrategy,
    /// Target sparsity ratio (0.0 to 1.0)
    pub sparsity_ratio: f64,
    /// Pruning granularity
    pub granularity: PruningGranularity,
    /// Optional pruning schedule for gradual pruning
    pub schedule: Option<PruningSchedule>,
    /// Random seed for deterministic pruning
    pub seed: Option<u64>,
}

/// Strategy for selecting weights to prune
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum PruningStrategy {
    /// Global magnitude pruning across all parameters
    Global,
    /// Layer-wise magnitude pruning within each layer
    LayerWise,
    /// Structured pruning (channels, filters)
    Structured {
        /// Type of structured pattern
        pattern_type: StructuredPattern,
    },
}

/// Granularity of pruning operations
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum PruningGranularity {
    /// Element-wise pruning (unstructured)
    Element,
    /// Channel-wise pruning
    Channel,
    /// Filter-wise pruning
    Filter,
    /// Block-wise pruning
    Block { height: usize, width: usize },
}

/// Structured pruning patterns
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum StructuredPattern {
    /// Channel pruning
    Channel,
    /// Filter pruning
    Filter,
    /// N:M sparsity pattern
    NM { n: usize, m: usize },
}

/// Pruning schedule for gradual sparsification during training
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum PruningSchedule {
    /// Linear schedule from initial to final sparsity
    Linear {
        /// Initial sparsity ratio
        initial_sparsity: f64,
        /// Final sparsity ratio  
        final_sparsity: f64,
        /// Duration in training steps
        duration_steps: usize,
    },
    /// Polynomial schedule
    Polynomial {
        /// Initial sparsity ratio
        initial_sparsity: f64,
        /// Final sparsity ratio
        final_sparsity: f64,
        /// Duration in training steps
        duration_steps: usize,
        /// Polynomial exponent
        exponent: f64,
    },
    /// Exponential schedule
    Exponential {
        /// Initial sparsity ratio
        initial_sparsity: f64,
        /// Final sparsity ratio
        final_sparsity: f64,
        /// Duration in training steps
        duration_steps: usize,
        /// Decay rate
        decay_rate: f64,
    },
}

/// Binary mask for indicating which weights to keep/prune
#[derive(Debug, Clone)]
pub struct PruningMask {
    /// Binary mask values (true = keep, false = prune)
    mask: Vec<bool>,
    /// Shape of the mask
    shape: Vec<usize>,
    /// Number of parameters kept
    kept_parameters: usize,
    /// Total number of parameters
    total_parameters: usize,
}

/// Statistics about pruning results
#[derive(Debug, Clone, PartialEq)]
pub struct PruningStats {
    /// Original parameter count
    pub original_params: usize,
    /// Parameters after pruning
    pub pruned_params: usize,
    /// Achieved sparsity ratio
    pub sparsity_ratio: f64,
    /// Compression ratio
    pub compression_ratio: f64,
    /// Memory savings (in bytes, estimated)
    pub memory_savings: usize,
}

/// Optimizer integration for pruned training
#[derive(Debug, Clone)]
pub struct PrunedOptimizer {
    /// Active pruning masks for each parameter group
    pub masks: HashMap<String, PruningMask>,
    /// Update count for scheduling
    pub update_count: usize,
    /// Pruning schedule configuration
    pub schedule_config: Option<PruningSchedule>,
}

/// Training state for pruned networks
#[derive(Debug, Clone)]
pub struct PrunedTrainingState {
    /// Current training step
    pub current_step: usize,
    /// Active sparsity ratio
    pub current_sparsity: f64,
    /// Accumulated pruning statistics
    pub stats: PruningStats,
    /// Whether gradual pruning is active
    pub gradual_pruning_active: bool,
}

/// Gradual pruning controller for training loops
#[derive(Debug)]
pub struct GradualPruningController {
    /// Pruning configuration
    config: PruningConfig,
    /// Current training state
    state: PrunedTrainingState,
    /// Parameter tensors being pruned
    parameters: HashMap<String, DenseTensor>,
    /// Current masks for each parameter
    masks: HashMap<String, PruningMask>,
}

/// Main struct for magnitude-based pruning operations
#[derive(Debug)]
pub struct MagnitudePruning;

/// Placeholder for tensor-like data structure
/// This will be replaced with actual Tensor integration once tensor API is stable
#[derive(Debug, Clone)]
pub struct DenseTensor {
    pub data: Vec<f32>,
    pub shape: Vec<usize>,
}

impl DenseTensor {
    pub fn new(data: Vec<f32>, shape: Vec<usize>) -> Self {
        assert_eq!(data.len(), shape.iter().product::<usize>());
        Self { data, shape }
    }
    
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }
    
    pub fn data(&self) -> &[f32] {
        &self.data
    }
    
    pub fn zeros(shape: Vec<usize>) -> Self {
        let total_elements = shape.iter().product::<usize>();
        Self {
            data: vec![0.0; total_elements],
            shape,
        }
    }
}

impl PruningConfig {
    /// Create global magnitude pruning configuration
    pub fn global(sparsity_ratio: f64) -> Self {
        Self {
            strategy: PruningStrategy::Global,
            sparsity_ratio,
            granularity: PruningGranularity::Element,
            schedule: None,
            seed: None,
        }
    }
    
    /// Create layer-wise magnitude pruning configuration
    pub fn layerwise(sparsity_ratio: f64) -> Self {
        Self {
            strategy: PruningStrategy::LayerWise,
            sparsity_ratio,
            granularity: PruningGranularity::Element,
            schedule: None,
            seed: None,
        }
    }
    
    /// Create structured magnitude pruning configuration
    pub fn structured(sparsity_ratio: f64, pattern: StructuredPattern) -> Self {
        let granularity = match pattern {
            StructuredPattern::Channel => PruningGranularity::Channel,
            StructuredPattern::Filter => PruningGranularity::Filter,
            StructuredPattern::NM { .. } => PruningGranularity::Element,
        };
        
        Self {
            strategy: PruningStrategy::Structured {
                pattern_type: pattern,
            },
            sparsity_ratio,
            granularity,
            schedule: None,
            seed: None,
        }
    }
    
    /// Add pruning schedule
    pub fn with_schedule(mut self, schedule: PruningSchedule) -> Self {
        self.schedule = Some(schedule);
        self
    }
    
    /// Set random seed for deterministic pruning
    pub fn with_seed(mut self, seed: u64) -> Self {
        self.seed = Some(seed);
        self
    }
    
    /// Validate configuration parameters
    pub fn validate(&self) -> Result<()> {
        // Validate sparsity ratio
        if !(0.0..=1.0).contains(&self.sparsity_ratio) {
            return Err(TransformerError::Config(
                "Sparsity ratio must be between 0.0 and 1.0".into()
            ));
        }
        
        // Validate schedule if present
        if let Some(ref schedule) = self.schedule {
            schedule.validate()?;
        }
        
        Ok(())
    }
}

impl PruningSchedule {
    /// Create linear pruning schedule
    pub fn linear(duration_steps: usize, final_sparsity: f64) -> Self {
        Self::Linear {
            initial_sparsity: 0.0,
            final_sparsity,
            duration_steps,
        }
    }
    
    /// Create polynomial pruning schedule
    pub fn polynomial(duration_steps: usize, exponent: f64) -> Self {
        Self::Polynomial {
            initial_sparsity: 0.0,
            final_sparsity: 0.9, // Default to 90% sparsity
            duration_steps,
            exponent,
        }
    }
    
    /// Create exponential pruning schedule
    pub fn exponential(duration_steps: usize, decay_rate: f64) -> Self {
        Self::Exponential {
            initial_sparsity: 0.0,
            final_sparsity: 0.9, // Default to 90% sparsity
            duration_steps,
            decay_rate,
        }
    }
    
    /// Compute sparsity ratio at given training step
    pub fn compute_sparsity(&self, current_step: usize) -> f64 {
        match self {
            Self::Linear { initial_sparsity, final_sparsity, duration_steps } => {
                if current_step >= *duration_steps {
                    return *final_sparsity;
                }
                let progress = current_step as f64 / *duration_steps as f64;
                initial_sparsity + progress * (final_sparsity - initial_sparsity)
            }
            Self::Polynomial { initial_sparsity, final_sparsity, duration_steps, exponent } => {
                if current_step >= *duration_steps {
                    return *final_sparsity;
                }
                let progress = current_step as f64 / *duration_steps as f64;
                let poly_progress = progress.powf(*exponent);
                initial_sparsity + poly_progress * (final_sparsity - initial_sparsity)
            }
            Self::Exponential { initial_sparsity, final_sparsity, duration_steps, decay_rate } => {
                if current_step >= *duration_steps {
                    return *final_sparsity;
                }
                let progress = current_step as f64 / *duration_steps as f64;
                let exp_progress = 1.0 - (-decay_rate * progress).exp();
                initial_sparsity + exp_progress * (final_sparsity - initial_sparsity)
            }
        }
    }
    
    /// Validate schedule parameters
    pub fn validate(&self) -> Result<()> {
        match self {
            Self::Linear { initial_sparsity, final_sparsity, duration_steps } => {
                Self::validate_common_params(*initial_sparsity, *final_sparsity, *duration_steps)?;
            }
            Self::Polynomial { initial_sparsity, final_sparsity, duration_steps, exponent } => {
                Self::validate_common_params(*initial_sparsity, *final_sparsity, *duration_steps)?;
                if *exponent <= 0.0 {
                    return Err(TransformerError::Config("Polynomial exponent must be positive".into()));
                }
            }
            Self::Exponential { initial_sparsity, final_sparsity, duration_steps, decay_rate } => {
                Self::validate_common_params(*initial_sparsity, *final_sparsity, *duration_steps)?;
                if *decay_rate <= 0.0 {
                    return Err(TransformerError::Config("Exponential decay rate must be positive".into()));
                }
            }
        }
        Ok(())
    }
    
    fn validate_common_params(initial: f64, final_val: f64, duration: usize) -> Result<()> {
        if !(0.0..=1.0).contains(&initial) {
            return Err(TransformerError::Config("Initial sparsity must be between 0.0 and 1.0".into()));
        }
        if !(0.0..=1.0).contains(&final_val) {
            return Err(TransformerError::Config("Final sparsity must be between 0.0 and 1.0".into()));
        }
        if duration == 0 {
            return Err(TransformerError::Config("Duration steps must be greater than 0".into()));
        }
        Ok(())
    }
}

impl PruningMask {
    /// Create new pruning mask
    pub fn new(mask: Vec<bool>, shape: Vec<usize>) -> Self {
        let total_parameters = shape.iter().product::<usize>();
        let kept_parameters = mask.iter().filter(|&&x| x).count();
        
        Self {
            mask,
            shape,
            kept_parameters,
            total_parameters,
        }
    }
    
    /// Get mask values
    pub fn mask(&self) -> &[bool] {
        &self.mask
    }
    
    /// Get mask shape
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }
    
    /// Get sparsity ratio
    pub fn sparsity_ratio(&self) -> f64 {
        1.0 - (self.kept_parameters as f64 / self.total_parameters as f64)
    }
    
    /// Apply mask to tensor (element-wise multiplication)
    pub fn apply(&self, tensor: &DenseTensor) -> Result<DenseTensor> {
        if tensor.shape != self.shape {
            return Err(TransformerError::Config("Tensor and mask shapes must match".into()));
        }
        
        let masked_data: Vec<f32> = tensor.data.iter()
            .zip(self.mask.iter())
            .map(|(&value, &keep)| if keep { value } else { 0.0 })
            .collect();
        
        Ok(DenseTensor::new(masked_data, tensor.shape.clone()))
    }
    
    /// Invert the mask (convert keep->prune and prune->keep)
    pub fn invert(&self) -> Self {
        let inverted_mask: Vec<bool> = self.mask.iter().map(|&x| !x).collect();
        Self::new(inverted_mask, self.shape.clone())
    }
    
    /// Apply mask to gradients (used in optimizer integration)
    pub fn apply_to_gradients(&self, gradients: &DenseTensor) -> Result<DenseTensor> {
        if gradients.shape != self.shape {
            return Err(TransformerError::Config("Gradient and mask shapes must match".into()));
        }
        
        // Zero out gradients for pruned weights
        let masked_gradients: Vec<f32> = gradients.data.iter()
            .zip(self.mask.iter())
            .map(|(&grad, &keep)| if keep { grad } else { 0.0 })
            .collect();
        
        Ok(DenseTensor::new(masked_gradients, gradients.shape.clone()))
    }
}

impl PruningStats {
    /// Create new pruning statistics
    pub fn new(original_params: usize, pruned_params: usize) -> Self {
        let sparsity_ratio = 1.0 - (pruned_params as f64 / original_params as f64);
        let compression_ratio = if pruned_params == 0 {
            f64::INFINITY
        } else {
            original_params as f64 / pruned_params as f64
        };
        
        // Estimate memory savings (assuming 4 bytes per f32)
        let memory_savings = (original_params - pruned_params) * 4;
        
        Self {
            original_params,
            pruned_params,
            sparsity_ratio,
            compression_ratio,
            memory_savings,
        }
    }
    
    /// Check if tensor meets target sparsity threshold
    pub fn meets_sparsity_target(&self, target_sparsity: f64) -> bool {
        self.sparsity_ratio >= target_sparsity
    }
    
    /// Merge statistics from multiple pruning operations
    pub fn merge(&self, other: &PruningStats) -> Self {
        let total_original = self.original_params + other.original_params;
        let total_pruned = self.pruned_params + other.pruned_params;
        Self::new(total_original, total_pruned)
    }
}

impl MagnitudePruning {
    /// Apply magnitude-based pruning to tensor
    ///
    /// # Arguments
    /// * `tensor` - Input tensor to prune
    /// * `config` - Pruning configuration
    ///
    /// # Returns
    /// Pruned tensor with smallest magnitude weights set to zero
    pub fn prune(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
        // Validate configuration
        config.validate()?;
        
        match config.strategy {
            PruningStrategy::Global => {
                Self::prune_global(tensor, config)
            }
            PruningStrategy::LayerWise => {
                Self::prune_layerwise(tensor, config)
            }
            PruningStrategy::Structured { ref pattern_type } => {
                Self::prune_structured(tensor, config, pattern_type)
            }
        }
    }
    
    /// Create pruning mask without modifying the original tensor
    ///
    /// # Arguments
    /// * `tensor` - Input tensor to analyze
    /// * `config` - Pruning configuration
    ///
    /// # Returns
    /// Binary mask indicating which weights to keep (true) or prune (false)
    pub fn create_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        // Validate configuration
        config.validate()?;
        
        match config.strategy {
            PruningStrategy::Global => {
                Self::create_global_mask(tensor, config)
            }
            PruningStrategy::LayerWise => {
                Self::create_layerwise_mask(tensor, config)
            }
            PruningStrategy::Structured { ref pattern_type } => {
                Self::create_structured_mask(tensor, config, pattern_type)
            }
        }
    }
    
    /// Compute importance scores based on magnitude
    ///
    /// # Arguments
    /// * `tensor` - Input tensor
    ///
    /// # Returns
    /// Vector of importance scores (absolute values)
    pub fn compute_importance_scores(tensor: &DenseTensor) -> Result<Vec<f32>> {
        let scores: Vec<f32> = tensor.data.iter().map(|&x| x.abs()).collect();
        Ok(scores)
    }
    
    /// Analyze pruning results
    ///
    /// # Arguments
    /// * `original` - Original tensor before pruning
    /// * `pruned` - Tensor after pruning
    ///
    /// # Returns
    /// Detailed statistics about the pruning results
    pub fn analyze_pruning(original: &DenseTensor, pruned: &DenseTensor) -> Result<PruningStats> {
        if original.shape != pruned.shape {
            return Err(TransformerError::Config("Tensor shapes must match for analysis".into()));
        }
        
        let original_nonzeros = original.data.iter().filter(|&&x| x != 0.0).count();
        let pruned_nonzeros = pruned.data.iter().filter(|&&x| x != 0.0).count();
        
        Ok(PruningStats::new(original_nonzeros, pruned_nonzeros))
    }
    
    // Private implementation methods - will be implemented in subsequent TDD iterations
    
    fn prune_global(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
        let mask = Self::create_global_mask(tensor, config)?;
        mask.apply(tensor)
    }
    
    fn prune_layerwise(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
        let mask = Self::create_layerwise_mask(tensor, config)?;
        mask.apply(tensor)
    }
    
    fn prune_structured(
        tensor: &DenseTensor,
        config: &PruningConfig,
        pattern: &StructuredPattern,
    ) -> Result<DenseTensor> {
        let mask = Self::create_structured_mask(tensor, config, pattern)?;
        mask.apply(tensor)
    }
    
    fn create_global_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        let total_elements = tensor.data.len();
        let elements_to_keep = ((1.0 - config.sparsity_ratio) * total_elements as f64).round() as usize;
        
        // Compute importance scores (absolute magnitudes)
        let importance_scores = Self::compute_importance_scores(tensor)?;
        
        // Create (index, score) pairs for sorting
        let mut indexed_scores: Vec<(usize, f32)> = importance_scores.iter()
            .enumerate()
            .map(|(i, &score)| (i, score))
            .collect();
        
        // Sort by importance score in descending order (keep highest magnitude values)
        indexed_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
        
        // Create mask: true for kept elements, false for pruned elements
        let mut mask_data = vec![false; total_elements];
        for i in 0..elements_to_keep {
            if i < indexed_scores.len() {
                let (index, _) = indexed_scores[i];
                mask_data[index] = true;
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
    
    fn create_layerwise_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        // For layer-wise pruning, we currently support 2D tensors (treating each row as a layer)
        if tensor.shape.len() != 2 {
            return Err(TransformerError::Config(
                "Layer-wise pruning currently only supports 2D tensors".into()
            ));
        }
        
        let rows = tensor.shape[0];
        let cols = tensor.shape[1];
        let elements_per_row = cols;
        let elements_to_keep_per_row = ((1.0 - config.sparsity_ratio) * elements_per_row as f64).round() as usize;
        
        let mut mask_data = vec![false; tensor.data.len()];
        
        // Process each row independently
        for row in 0..rows {
            let row_start = row * cols;
            let row_end = row_start + cols;
            
            // Get magnitudes for this row
            let row_scores: Vec<(usize, f32)> = (row_start..row_end)
                .map(|i| (i, tensor.data[i].abs()))
                .collect();
            
            // Sort by magnitude in descending order
            let mut sorted_scores = row_scores;
            sorted_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
            
            // Keep the top elements_to_keep_per_row elements in this row
            for i in 0..elements_to_keep_per_row.min(sorted_scores.len()) {
                let (original_index, _) = sorted_scores[i];
                mask_data[original_index] = true;
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
    
    fn create_structured_mask(
        tensor: &DenseTensor,
        config: &PruningConfig,
        pattern: &StructuredPattern,
    ) -> Result<PruningMask> {
        match pattern {
            StructuredPattern::Channel => {
                Self::create_channel_pruning_mask(tensor, config)
            }
            StructuredPattern::Filter => {
                Self::create_filter_pruning_mask(tensor, config)
            }
            StructuredPattern::NM { n, m } => {
                Self::create_nm_pruning_mask(tensor, config, *n, *m)
            }
        }
    }
    
    fn create_channel_pruning_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        // Channel pruning for 2D tensors (each row is a channel)
        if tensor.shape.len() != 2 {
            return Err(TransformerError::Config(
                "Channel pruning currently only supports 2D tensors".into()
            ));
        }
        
        let rows = tensor.shape[0];
        let cols = tensor.shape[1];
        let channels_to_keep = ((1.0 - config.sparsity_ratio) * rows as f64).round() as usize;
        
        // Calculate channel importance (sum of absolute values per row)
        let mut channel_scores: Vec<(usize, f32)> = Vec::new();
        
        for row in 0..rows {
            let row_start = row * cols;
            let row_end = row_start + cols;
            let row_sum: f32 = tensor.data[row_start..row_end]
                .iter()
                .map(|&x| x.abs())
                .sum();
            channel_scores.push((row, row_sum));
        }
        
        // Sort by importance in descending order
        channel_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
        
        // Create mask: keep top channels, zero out others
        let mut mask_data = vec![false; tensor.data.len()];
        for i in 0..channels_to_keep {
            if i < channel_scores.len() {
                let (channel_idx, _) = channel_scores[i];
                let row_start = channel_idx * cols;
                let row_end = row_start + cols;
                for j in row_start..row_end {
                    mask_data[j] = true;
                }
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
    
    fn create_filter_pruning_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        // Filter pruning for 2D tensors (each column is a filter)
        if tensor.shape.len() != 2 {
            return Err(TransformerError::Config(
                "Filter pruning currently only supports 2D tensors".into()
            ));
        }
        
        let rows = tensor.shape[0];
        let cols = tensor.shape[1];
        let filters_to_keep = ((1.0 - config.sparsity_ratio) * cols as f64).round() as usize;
        
        // Calculate filter importance (sum of absolute values per column)
        let mut filter_scores: Vec<(usize, f32)> = Vec::new();
        
        for col in 0..cols {
            let col_sum: f32 = (0..rows)
                .map(|row| tensor.data[row * cols + col].abs())
                .sum();
            filter_scores.push((col, col_sum));
        }
        
        // Sort by importance in descending order
        filter_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
        
        // Create mask: keep top filters, zero out others
        let mut mask_data = vec![false; tensor.data.len()];
        for i in 0..filters_to_keep {
            if i < filter_scores.len() {
                let (filter_idx, _) = filter_scores[i];
                for row in 0..rows {
                    mask_data[row * cols + filter_idx] = true;
                }
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
    
    fn create_nm_pruning_mask(tensor: &DenseTensor, config: &PruningConfig, n: usize, m: usize) -> Result<PruningMask> {
        // N:M sparsity: keep N elements out of every M consecutive elements
        if n == 0 || m == 0 {
            return Err(TransformerError::Config("N and M must be non-zero for N:M sparsity".into()));
        }
        if n > m {
            return Err(TransformerError::Config("N cannot be greater than M in N:M sparsity".into()));
        }
        
        let total_elements = tensor.data.len();
        let mut mask_data = vec![false; total_elements];
        
        // Process in groups of M elements
        for group_start in (0..total_elements).step_by(m) {
            let group_end = (group_start + m).min(total_elements);
            let group_size = group_end - group_start;
            let elements_to_keep = n.min(group_size);
            
            // Get magnitudes for this group
            let mut group_scores: Vec<(usize, f32)> = (group_start..group_end)
                .map(|i| (i, tensor.data[i].abs()))
                .collect();
            
            // Sort by magnitude in descending order
            group_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
            
            // Keep top N elements in this group
            for i in 0..elements_to_keep {
                let (original_index, _) = group_scores[i];
                mask_data[original_index] = true;
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
}

impl PrunedOptimizer {
    /// Create new pruned optimizer state
    pub fn new() -> Self {
        Self {
            masks: HashMap::new(),
            update_count: 0,
            schedule_config: None,
        }
    }
    
    /// Create pruned optimizer with schedule
    pub fn with_schedule(schedule: PruningSchedule) -> Self {
        Self {
            masks: HashMap::new(),
            update_count: 0,
            schedule_config: Some(schedule),
        }
    }
    
    /// Add parameter mask for tracking
    pub fn add_parameter_mask(&mut self, name: &str, mask: PruningMask) {
        self.masks.insert(name.to_string(), mask);
    }
    
    /// Get current sparsity ratio based on schedule
    pub fn get_current_sparsity(&self) -> f64 {
        if let Some(ref schedule) = self.schedule_config {
            schedule.compute_sparsity(self.update_count)
        } else {
            // If no schedule, return average sparsity of all masks
            if self.masks.is_empty() {
                0.0
            } else {
                let total_sparsity: f64 = self.masks.values()
                    .map(|mask| mask.sparsity_ratio())
                    .sum();
                total_sparsity / self.masks.len() as f64
            }
        }
    }
    
    /// Apply pruning masks to gradients before optimizer step
    pub fn apply_masks_to_gradients(&self, gradients: &HashMap<String, DenseTensor>) -> Result<HashMap<String, DenseTensor>> {
        let mut masked_gradients = HashMap::new();
        
        for (param_name, grad_tensor) in gradients {
            if let Some(mask) = self.masks.get(param_name) {
                let masked_grad = mask.apply_to_gradients(grad_tensor)?;
                masked_gradients.insert(param_name.clone(), masked_grad);
            } else {
                // No mask for this parameter, keep gradients as-is
                masked_gradients.insert(param_name.clone(), grad_tensor.clone());
            }
        }
        
        Ok(masked_gradients)
    }
    
    /// Update masks based on current parameters and schedule
    pub fn update_masks(&mut self, parameters: &HashMap<String, DenseTensor>, config: &PruningConfig) -> Result<()> {
        let target_sparsity = if let Some(ref schedule) = self.schedule_config {
            schedule.compute_sparsity(self.update_count)
        } else {
            config.sparsity_ratio
        };
        
        // Create updated config with current target sparsity
        let mut updated_config = config.clone();
        updated_config.sparsity_ratio = target_sparsity;
        
        // Update masks for each parameter
        for (param_name, tensor) in parameters {
            let new_mask = MagnitudePruning::create_mask(tensor, &updated_config)?;
            self.masks.insert(param_name.clone(), new_mask);
        }
        
        self.update_count += 1;
        Ok(())
    }
    
    /// Get pruning statistics across all parameters
    pub fn get_overall_stats(&self, parameters: &HashMap<String, DenseTensor>) -> Result<PruningStats> {
        let mut total_original = 0;
        let mut total_pruned = 0;
        
        for (param_name, tensor) in parameters {
            let original_count = tensor.data.len();
            total_original += original_count;
            
            if let Some(mask) = self.masks.get(param_name) {
                total_pruned += mask.kept_parameters;
            } else {
                // No mask means all parameters are kept
                total_pruned += original_count;
            }
        }
        
        Ok(PruningStats::new(total_original, total_pruned))
    }
}

impl PrunedTrainingState {
    /// Create new training state
    pub fn new() -> Self {
        Self {
            current_step: 0,
            current_sparsity: 0.0,
            stats: PruningStats::new(0, 0),
            gradual_pruning_active: false,
        }
    }
    
    /// Update training state
    pub fn update_step(&mut self, new_stats: PruningStats, sparsity: f64) {
        self.current_step += 1;
        self.current_sparsity = sparsity;
        self.stats = new_stats;
    }
    
    /// Check if gradual pruning should be applied
    pub fn should_update_pruning(&self, update_frequency: usize) -> bool {
        self.gradual_pruning_active && (self.current_step % update_frequency == 0)
    }
}

impl GradualPruningController {
    /// Create new gradual pruning controller
    pub fn new(config: PruningConfig) -> Self {
        Self {
            config,
            state: PrunedTrainingState::new(),
            parameters: HashMap::new(),
            masks: HashMap::new(),
        }
    }
    
    /// Add parameter to be pruned
    pub fn add_parameter(&mut self, name: &str, tensor: DenseTensor) {
        self.parameters.insert(name.to_string(), tensor);
    }
    
    /// Update parameters and recompute masks if needed
    pub fn update_parameters(&mut self, parameters: &HashMap<String, DenseTensor>) -> Result<()> {
        // Update stored parameters
        for (name, tensor) in parameters {
            self.parameters.insert(name.clone(), tensor.clone());
        }
        
        // Compute current sparsity based on schedule
        let current_sparsity = if let Some(ref schedule) = self.config.schedule {
            schedule.compute_sparsity(self.state.current_step)
        } else {
            self.config.sparsity_ratio
        };
        
        // Update config with current sparsity
        let mut current_config = self.config.clone();
        current_config.sparsity_ratio = current_sparsity;
        
        // Recompute masks for all parameters
        for (name, tensor) in &self.parameters {
            let mask = MagnitudePruning::create_mask(tensor, &current_config)?;
            self.masks.insert(name.clone(), mask);
        }
        
        // Update training state
        let overall_stats = self.compute_overall_stats();
        self.state.update_step(overall_stats, current_sparsity);
        
        Ok(())
    }
    
    /// Get current pruning masks
    pub fn get_masks(&self) -> &HashMap<String, PruningMask> {
        &self.masks
    }
    
    /// Get current training state
    pub fn get_state(&self) -> &PrunedTrainingState {
        &self.state
    }
    
    /// Apply current masks to parameters
    pub fn apply_pruning(&self, parameters: &HashMap<String, DenseTensor>) -> Result<HashMap<String, DenseTensor>> {
        let mut pruned_parameters = HashMap::new();
        
        for (name, tensor) in parameters {
            if let Some(mask) = self.masks.get(name) {
                let pruned_tensor = mask.apply(tensor)?;
                pruned_parameters.insert(name.clone(), pruned_tensor);
            } else {
                // No mask, keep original
                pruned_parameters.insert(name.clone(), tensor.clone());
            }
        }
        
        Ok(pruned_parameters)
    }
    
    /// Enable gradual pruning mode
    pub fn enable_gradual_pruning(&mut self) {
        self.state.gradual_pruning_active = true;
    }
    
    /// Disable gradual pruning mode
    pub fn disable_gradual_pruning(&mut self) {
        self.state.gradual_pruning_active = false;
    }
    
    /// Compute overall statistics
    fn compute_overall_stats(&self) -> PruningStats {
        let mut total_original = 0;
        let mut total_kept = 0;
        
        for (name, tensor) in &self.parameters {
            total_original += tensor.data.len();
            
            if let Some(mask) = self.masks.get(name) {
                total_kept += mask.kept_parameters;
            } else {
                total_kept += tensor.data.len();
            }
        }
        
        PruningStats::new(total_original, total_kept)
    }
}

#[cfg(all(test, feature = "disabled_tests"))]
mod magnitude_pruning_tests {
    use super::*;
    
    // Tests will be implemented following TDD principles
}