//! Pruning module for model compression.

use forge_shared::{ModelInfo, PruneMethod, PruneSchedule, PruningConfig};

/// Pruner for model weight pruning.
#[derive(Debug)]
pub struct Pruner {
    /// Current pruning step.
    current_step: usize,
    /// RNG state.
    rng_state: u64,
}

impl Default for Pruner {
    fn default() -> Self {
        Self::new()
    }
}

impl Pruner {
    /// Create a new pruner.
    pub fn new() -> Self {
        Self {
            current_step: 0,
            rng_state: 42,
        }
    }

    /// Prune a model and return the fraction of parameters remaining.
    pub fn prune(&mut self, model: &ModelInfo, config: &PruningConfig) -> f32 {
        // Calculate effective sparsity
        let _target_sparsity = config.target_sparsity;

        // Apply schedule
        let effective_sparsity = self.apply_schedule(config);

        // Adjust for method characteristics
        let method_factor = self.method_efficiency(config.method);

        // Calculate remaining parameters
        let remaining = 1.0 - effective_sparsity * method_factor;

        // Account for excluded layers
        let exclude_fraction = config.exclude_layers.len() as f32 / model.num_layers.max(1) as f32;
        let adjusted_remaining = remaining * (1.0 - exclude_fraction) + exclude_fraction;

        adjusted_remaining.clamp(0.1, 1.0)
    }

    /// Apply pruning schedule to get current sparsity.
    fn apply_schedule(&mut self, config: &PruningConfig) -> f32 {
        self.current_step += 1;
        let progress = (self.current_step as f32 / config.num_steps.max(1) as f32).min(1.0);

        match config.schedule {
            PruneSchedule::OneShot => config.final_sparsity,
            PruneSchedule::Linear => {
                config.initial_sparsity
                    + progress * (config.final_sparsity - config.initial_sparsity)
            }
            PruneSchedule::Gradual => {
                // Gradual: slower start, faster end
                let t = progress * progress;
                config.initial_sparsity + t * (config.final_sparsity - config.initial_sparsity)
            }
            PruneSchedule::Cubic => {
                let t = progress * progress * progress;
                config.initial_sparsity + t * (config.final_sparsity - config.initial_sparsity)
            }
            PruneSchedule::Exponential => {
                let t = 1.0 - (-3.0 * progress).exp();
                config.initial_sparsity + t * (config.final_sparsity - config.initial_sparsity)
            }
        }
    }

    /// Get efficiency factor for pruning method.
    fn method_efficiency(&self, method: PruneMethod) -> f32 {
        match method {
            PruneMethod::Magnitude => 0.95, // Simple, may miss important weights
            PruneMethod::Movement => 0.98,  // Considers weight change during training
            PruneMethod::Structured => 1.0, // Full efficiency for structured
            PruneMethod::Unstructured => 0.9, // Sparse storage overhead
            PruneMethod::LayerWise => 0.95,
            PruneMethod::Global => 0.97,
            PruneMethod::LotteryTicket => 0.99, // Very effective
            PruneMethod::NM => 1.0,             // Hardware efficient
        }
    }

    /// Estimate accuracy degradation from pruning.
    pub fn estimate_degradation(&self, config: &PruningConfig) -> f64 {
        let sparsity = config.target_sparsity as f64;

        // Degradation increases with sparsity
        let base_degradation = sparsity * sparsity * 0.1;

        // Method factor
        let method_factor = match config.method {
            PruneMethod::Magnitude => 1.2,
            PruneMethod::Movement => 0.8,
            PruneMethod::LotteryTicket => 0.6,
            PruneMethod::Structured => 1.5, // Structured is more aggressive
            _ => 1.0,
        };

        // Retraining helps recover accuracy
        let retrain_factor = 1.0 / (1.0 + config.retrain_epochs as f64 * 0.1);

        base_degradation * method_factor * retrain_factor
    }

    /// Calculate weight importance scores.
    pub fn compute_importance(&mut self, num_weights: usize) -> Vec<f64> {
        let mut importance = Vec::with_capacity(num_weights);

        for _ in 0..num_weights {
            // Simulate importance scores (would be actual weight magnitudes in practice)
            importance.push(self.random().abs());
        }

        importance
    }

    /// Generate pruning mask based on importance.
    pub fn generate_mask(&mut self, importance: &[f64], sparsity: f32) -> Vec<bool> {
        let n = importance.len();
        let num_prune = (n as f32 * sparsity) as usize;

        // Find threshold
        let mut sorted = importance.to_vec();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let threshold = if num_prune < n {
            sorted[num_prune]
        } else {
            f64::MAX
        };

        // Generate mask (true = keep, false = prune)
        importance.iter().map(|&imp| imp >= threshold).collect()
    }

    /// Check if weight pattern matches N:M sparsity.
    pub fn check_nm_pattern(&self, mask: &[bool], n: usize, m: usize) -> bool {
        if !mask.len().is_multiple_of(m) {
            return false;
        }

        for chunk in mask.chunks(m) {
            let active_count = chunk.iter().filter(|&&x| x).count();
            if active_count != n {
                return false;
            }
        }

        true
    }

    /// Random number.
    fn random(&mut self) -> f64 {
        self.rng_state = self
            .rng_state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        (self.rng_state >> 11) as f64 / (1u64 << 53) as f64
    }

    /// Reset pruner state.
    pub fn reset(&mut self) {
        self.current_step = 0;
    }
}

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

    #[test]
    fn test_pruner_creation() {
        let pruner = Pruner::new();
        assert_eq!(pruner.current_step, 0);
    }

    #[test]
    fn test_prune_50_percent() {
        let mut pruner = Pruner::new();
        let model = forge_shared::sample_model_info();
        let config = PruningConfig {
            target_sparsity: 0.5,
            final_sparsity: 0.5,
            schedule: PruneSchedule::OneShot,
            ..Default::default()
        };

        let remaining = pruner.prune(&model, &config);
        assert!(remaining > 0.4 && remaining < 0.6);
    }

    #[test]
    fn test_gradual_schedule() {
        let mut pruner = Pruner::new();
        let config = PruningConfig {
            initial_sparsity: 0.0,
            final_sparsity: 0.9,
            schedule: PruneSchedule::Gradual,
            num_steps: 10,
            ..Default::default()
        };

        let mut prev_sparsity = 0.0;
        for _ in 0..10 {
            let sparsity = pruner.apply_schedule(&config);
            assert!(sparsity >= prev_sparsity);
            prev_sparsity = sparsity;
        }
    }

    #[test]
    fn test_estimate_degradation() {
        let pruner = Pruner::new();

        let config_low = PruningConfig {
            target_sparsity: 0.3,
            ..Default::default()
        };
        let config_high = PruningConfig {
            target_sparsity: 0.9,
            ..Default::default()
        };

        let deg_low = pruner.estimate_degradation(&config_low);
        let deg_high = pruner.estimate_degradation(&config_high);

        // Higher sparsity should have higher degradation
        assert!(deg_high > deg_low);
    }

    #[test]
    fn test_compute_importance() {
        let mut pruner = Pruner::new();
        let importance = pruner.compute_importance(100);

        assert_eq!(importance.len(), 100);
        for &imp in &importance {
            assert!(imp >= 0.0 && imp <= 1.0);
        }
    }

    #[test]
    fn test_generate_mask() {
        let mut pruner = Pruner::new();
        let importance = vec![0.1, 0.5, 0.2, 0.8, 0.3, 0.9, 0.4, 0.7, 0.6, 0.05];

        let mask = pruner.generate_mask(&importance, 0.5);

        assert_eq!(mask.len(), 10);
        let kept = mask.iter().filter(|&&x| x).count();
        assert_eq!(kept, 5); // 50% kept
    }

    #[test]
    fn test_nm_pattern() {
        let pruner = Pruner::new();

        // Valid 2:4 pattern
        let mask_valid = vec![true, true, false, false, true, false, true, false];
        assert!(pruner.check_nm_pattern(&mask_valid, 2, 4));

        // Invalid pattern
        let mask_invalid = vec![true, true, true, false];
        assert!(!pruner.check_nm_pattern(&mask_invalid, 2, 4));
    }

    #[test]
    fn test_reset() {
        let mut pruner = Pruner::new();
        pruner.current_step = 5;
        pruner.reset();
        assert_eq!(pruner.current_step, 0);
    }
}