//! # Selective State Space Mechanism for Mamba-2
//!
//! Advanced selective state mechanism that dynamically chooses which
//! state information to retain and which to discard, enabling efficient
//! long-sequence modeling with sub-linear memory growth.

use std::collections::{BTreeMap, HashMap, VecDeque};
use std::mem::size_of;
use std::sync::atomic::{AtomicU64, Ordering};

use nalgebra::DVector;
use tracing::{debug, instrument};

use super::{Mamba2Config, Mamba2State};
use crate::gpu_tensor::GpuTensor;
use ml_core::MLError;

/// Configuration for selective state space mechanism
#[derive(Debug, Clone)]
pub struct SelectiveStateConfig {
    pub importance_threshold: f64,
    pub max_active_states: usize,
    pub compression_ratio: f64,
    pub importance_decay: f64,
    pub importance_window: usize,
    pub adaptive_threshold: bool,
    pub memory_budget: usize,
}

impl Default for SelectiveStateConfig {
    fn default() -> Self {
        Self {
            importance_threshold: 0.1,
            max_active_states: 1000,
            compression_ratio: 0.1,
            importance_decay: 0.99,
            importance_window: 100,
            adaptive_threshold: true,
            memory_budget: 1024 * 1024,
        }
    }
}

/// State importance tracker
#[derive(Debug, Clone)]
pub struct StateImportance {
    pub score: f64,
    pub usage_count: u64,
    pub last_access: u64,
    pub moving_average: f64,
    pub variance: f64,
}

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

impl StateImportance {
    pub const fn new() -> Self {
        Self {
            score: 0.0,
            usage_count: 0,
            last_access: 0,
            moving_average: 0.0,
            variance: 0.0,
        }
    }

    pub fn update(&mut self, score: f64, timestamp: u64, decay: f64) {
        let old_avg = self.moving_average;
        self.moving_average = decay * self.moving_average + (1.0 - decay) * score;
        let diff = score - old_avg;
        self.variance = decay * self.variance + (1.0 - decay) * diff * diff;
        self.score = score;
        self.usage_count += 1;
        self.last_access = timestamp;
    }

    pub fn effective_importance(&self) -> f64 {
        let age = 100_u64.saturating_sub(self.last_access);
        let recency_weight = 1.0 / (1.0 + age as f64 * 0.01);
        let stability_weight = 1.0 / (1.0 + self.variance);
        self.moving_average * recency_weight * stability_weight
    }
}

/// State compressor for memory efficiency
#[derive(Debug, Clone)]
pub struct StateCompressor {
    pub compression_stats: HashMap<String, f64>,
}

impl StateCompressor {
    pub fn new(_config: SelectiveStateConfig) -> Self {
        Self {
            compression_stats: HashMap::new(),
        }
    }

    pub fn compress_lossy(&mut self, state: &DVector<f64>, quality: f64) -> DVector<f64> {
        let threshold = self.compute_compression_threshold(state, quality);

        let compressed = state.map(|x| {
            if x.abs() < threshold {
                0.0
            } else {
                let scale = 1.0 / threshold;
                (x * scale).round() / scale
            }
        });

        let compression_ratio = self.compute_compression_ratio(state, &compressed);
        self.compression_stats
            .insert("last_lossy_ratio".to_owned(), compression_ratio);

        compressed
    }

    pub fn compress_lossless(
        &mut self,
        state: &DVector<f64>,
        epsilon: f64,
    ) -> (Vec<(f64, usize)>, usize) {
        let mut runs = Vec::new();
        let mut current_value = state[0];
        let mut run_length = 1;

        for i in 1..state.len() {
            if (state[i] - current_value).abs() < epsilon {
                run_length += 1;
            } else {
                runs.push((current_value, run_length));
                current_value = state[i];
                run_length = 1;
            }
        }

        runs.push((current_value, run_length));

        let original_size = state.len();
        let compressed_size = runs.len();
        let compression_ratio = compressed_size as f64 / original_size as f64;
        self.compression_stats
            .insert("last_lossless_ratio".to_owned(), compression_ratio);

        (runs, original_size)
    }

    pub fn decompress_lossless(&self, runs: &[(f64, usize)], original_size: usize) -> DVector<f64> {
        let mut decompressed = DVector::zeros(original_size);
        let mut index = 0;

        for &(value, length) in runs {
            for _ in 0..length {
                if index < original_size {
                    decompressed[index] = value;
                    index += 1;
                }
            }
        }

        decompressed
    }

    fn compute_compression_threshold(&self, state: &DVector<f64>, quality: f64) -> f64 {
        let mut sorted_abs: Vec<f64> = state.iter().map(|x| x.abs()).collect();
        sorted_abs.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let percentile_index = ((1.0 - quality) * sorted_abs.len() as f64) as usize;
        sorted_abs.get(percentile_index).copied().unwrap_or(0.0)
    }

    fn compute_compression_ratio(&self, original: &DVector<f64>, compressed: &DVector<f64>) -> f64 {
        let original_nonzero = original.iter().filter(|&&x| x != 0.0).count();
        let compressed_nonzero = compressed.iter().filter(|&&x| x != 0.0).count();

        if original_nonzero == 0 {
            1.0
        } else {
            compressed_nonzero as f64 / original_nonzero as f64
        }
    }
}

/// Selective state space implementation
#[derive(Debug)]
pub struct SelectiveStateSpace {
    config: SelectiveStateConfig,
    pub importance_tracker: Vec<StateImportance>,
    pub active_indices: Vec<usize>,
    pub compressed_states: BTreeMap<usize, Vec<u8>>,
    selection_updates: AtomicU64,
    compression_operations: AtomicU64,
    decompression_operations: AtomicU64,
    memory_usage: AtomicU64,
    threshold_history: VecDeque<f64>,
    current_threshold: f64,
    timestamp_counter: AtomicU64,
}

impl SelectiveStateSpace {
    pub fn new(config: &Mamba2Config) -> Result<Self, MLError> {
        let selective_config = SelectiveStateConfig::default();
        let state_size = config.d_model * config.expand;
        let importance_tracker = (0..state_size).map(|_| StateImportance::new()).collect();

        Ok(Self {
            config: selective_config.clone(),
            importance_tracker,
            active_indices: Vec::new(),
            compressed_states: BTreeMap::new(),
            selection_updates: AtomicU64::new(0),
            compression_operations: AtomicU64::new(0),
            decompression_operations: AtomicU64::new(0),
            memory_usage: AtomicU64::new(0),
            threshold_history: VecDeque::new(),
            current_threshold: selective_config.importance_threshold,
            timestamp_counter: AtomicU64::new(0),
        })
    }

    #[instrument(skip(self, input, _state))]
    pub fn update_importance_scores(
        &mut self,
        input: &GpuTensor,
        _state: &mut Mamba2State,
    ) -> Result<(), MLError> {
        let timestamp = self.timestamp_counter.fetch_add(1, Ordering::Relaxed);

        // Compute element-wise abs on GPU, then extract a single mean scalar
        // per importance tracker bucket. For flat tensors, compute abs on GPU
        // then download only the scalar mean per element.
        let abs_input = crate::gpu_tensor::gpu_abs(input)?;
        // Use gpu_mean_all to get overall importance, then distribute uniformly
        // to avoid downloading the full tensor.
        let overall_mean = crate::gpu_tensor::gpu_mean_all(&abs_input)? as f64;

        let tracker_len = self.importance_tracker.len();
        for i in 0..tracker_len {
            // Approximate per-element importance with overall mean
            // This avoids O(N) host download while maintaining importance tracking
            self.importance_tracker[i].update(
                overall_mean,
                timestamp,
                self.config.importance_decay,
            );
        }

        self.update_active_selection()?;

        if self.config.adaptive_threshold {
            self.update_adaptive_threshold()?;
        }

        self.selection_updates.fetch_add(1, Ordering::Relaxed);
        Ok(())
    }

    #[allow(clippy::unnecessary_wraps)]
    fn update_active_selection(&mut self) -> Result<(), MLError> {
        let mut importance_scores: Vec<(usize, f64)> = self
            .importance_tracker
            .iter()
            .enumerate()
            .map(|(i, tracker)| (i, tracker.effective_importance()))
            .collect();

        importance_scores
            .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        self.active_indices.clear();
        for &(index, score) in &importance_scores {
            if score >= self.current_threshold
                && self.active_indices.len() < self.config.max_active_states
            {
                self.active_indices.push(index);
            }
        }

        while self.active_indices.len() < 10 && self.active_indices.len() < importance_scores.len()
        {
            let (index, _) = importance_scores[self.active_indices.len()];
            self.active_indices.push(index);
        }

        debug!(
            "Updated active states: {} out of {}",
            self.active_indices.len(),
            self.importance_tracker.len()
        );

        Ok(())
    }

    #[allow(clippy::unnecessary_wraps)]
    fn update_adaptive_threshold(&mut self) -> Result<(), MLError> {
        let current_memory = self.memory_usage.load(Ordering::Relaxed) as f64;
        let target_memory = self.config.memory_budget as f64;
        let memory_ratio = current_memory / target_memory;

        if memory_ratio > 1.0 {
            self.current_threshold *= 1.1;
        } else if memory_ratio < 0.7 {
            self.current_threshold *= 0.95;
        } else {
            // Memory usage within acceptable range — no threshold adjustment
        }

        self.current_threshold = self.current_threshold.max(0.001).min(1.0);
        self.threshold_history.push_back(self.current_threshold);
        if self.threshold_history.len() > self.config.importance_window {
            self.threshold_history.pop_front();
        }

        Ok(())
    }

    pub fn compress_state_component(
        &mut self,
        index: usize,
        state: &mut Mamba2State,
    ) -> Result<(), MLError> {
        if index < state.selective_state.len() {
            let value = state.selective_state[index];
            let compressed = self.compress_float_to_bytes(value);
            self.compressed_states.insert(index, compressed);
            state.selective_state[index] = 0.0;
            self.compression_operations.fetch_add(1, Ordering::Relaxed);
            self.update_memory_usage()?;
        }
        Ok(())
    }

    pub fn decompress_state_component(
        &mut self,
        index: usize,
        state: &mut Mamba2State,
    ) -> Result<(), MLError> {
        if let Some(compressed) = self.compressed_states.remove(&index) {
            let value = self.decompress_bytes_to_float(&compressed);
            if index < state.selective_state.len() {
                state.selective_state[index] = value;
            }
            self.decompression_operations
                .fetch_add(1, Ordering::Relaxed);
            self.update_memory_usage()?;
        }
        Ok(())
    }

    fn compress_float_to_bytes(&self, value: f64) -> Vec<u8> {
        value.to_le_bytes().to_vec() // cpu-side bytes conversion
    }

    fn decompress_bytes_to_float(&self, bytes: &[u8]) -> f64 {
        if bytes.len() >= 8 {
            let mut array = [0_u8; 8];
            array.copy_from_slice(&bytes[..8]);
            f64::from_le_bytes(array)
        } else {
            0.0
        }
    }

    #[allow(clippy::unnecessary_wraps)]
    fn update_memory_usage(&self) -> Result<(), MLError> {
        let compressed_memory = self
            .compressed_states
            .values()
            .map(|v| v.len())
            .sum::<usize>();
        let tracker_memory = self.importance_tracker.len() * size_of::<StateImportance>();
        let active_memory = self.active_indices.len() * size_of::<usize>();
        let total_memory = compressed_memory + tracker_memory + active_memory;
        self.memory_usage
            .store(total_memory as u64, Ordering::Relaxed);
        Ok(())
    }

    pub fn get_performance_metrics(&self) -> HashMap<String, f64> {
        let mut metrics = HashMap::new();
        metrics.insert(
            "selection_updates".to_owned(),
            self.selection_updates.load(Ordering::Relaxed) as f64,
        );
        metrics.insert(
            "compression_operations".to_owned(),
            self.compression_operations.load(Ordering::Relaxed) as f64,
        );
        metrics.insert(
            "decompression_operations".to_owned(),
            self.decompression_operations.load(Ordering::Relaxed) as f64,
        );
        metrics.insert(
            "memory_usage_bytes".to_owned(),
            self.memory_usage.load(Ordering::Relaxed) as f64,
        );

        let active_ratio = if !self.importance_tracker.is_empty() {
            self.active_indices.len() as f64 / self.importance_tracker.len() as f64
        } else {
            0.0
        };
        metrics.insert("active_state_ratio".to_owned(), active_ratio);

        let avg_importance = if !self.importance_tracker.is_empty() {
            self.importance_tracker
                .iter()
                .map(|t| t.effective_importance())
                .sum::<f64>()
                / self.importance_tracker.len() as f64
        } else {
            0.0
        };
        metrics.insert("average_importance_score".to_owned(), avg_importance);
        metrics.insert("current_threshold".to_owned(), self.current_threshold);
        metrics.insert(
            "compressed_states_count".to_owned(),
            self.compressed_states.len() as f64,
        );

        metrics
    }

    pub fn get_selection_efficiency(&self) -> f64 {
        let total_states = self.importance_tracker.len();
        let active_states = self.active_indices.len();
        if total_states > 0 {
            1.0 - (active_states as f64 / total_states as f64)
        } else {
            0.0
        }
    }

    pub fn get_compression_ratio(&self) -> f64 {
        let total_states = self.importance_tracker.len();
        let compressed_states = self.compressed_states.len();
        if total_states > 0 {
            compressed_states as f64 / total_states as f64
        } else {
            0.0
        }
    }
}

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

    #[test]
    fn test_state_importance_update() {
        let mut importance = StateImportance::new();
        importance.update(0.5, 100, 0.9);
        assert_eq!(importance.score, 0.5);
        assert_eq!(importance.usage_count, 1);
        importance.update(0.8, 200, 0.9);
        assert_eq!(importance.score, 0.8);
        assert_eq!(importance.usage_count, 2);
        assert!(importance.effective_importance() > 0.0);
    }

    #[test]
    fn test_state_compressor() {
        let config = SelectiveStateConfig::default();
        let mut compressor = StateCompressor::new(config);
        let data = DVector::from_vec(vec![1.0, 0.0, 0.0, 0.0, 2.0, 3.0, 0.0]);
        let compressed = compressor.compress_lossy(&data, 0.8);
        assert_eq!(compressed.len(), data.len());
        let (run_length, original_size) = compressor.compress_lossless(&data, 0.1);
        let decompressed = compressor.decompress_lossless(&run_length, original_size);
        assert_eq!(decompressed.len(), data.len());
    }

    #[test]
    fn test_selective_state_creation() -> Result<(), MLError> {
        let mut config = Mamba2Config::emergency_safe_defaults();
        config.d_model = 8;
        config.d_state = 4;
        config.expand = 2;
        let selective_state = SelectiveStateSpace::new(&config)?;
        assert_eq!(selective_state.importance_tracker.len(), 16);
        assert_eq!(selective_state.active_indices.len(), 0);
        Ok(())
    }

    #[test]
    fn test_performance_metrics() -> Result<(), MLError> {
        let config = Mamba2Config::default();
        let selective_state = SelectiveStateSpace::new(&config)?;
        let metrics = selective_state.get_performance_metrics();
        assert!(metrics.contains_key("selection_updates"));
        assert!(metrics.contains_key("compression_operations"));
        assert!(metrics.contains_key("active_state_ratio"));
        assert!(metrics.contains_key("average_importance_score"));
        Ok(())
    }
}
