compression_prompt/

statistical_filter.rs

//! Statistical token importance filtering (LLMLingua-inspired, model-free)
//!
//! This module implements a compression strategy similar to LLMLingua but using
//! pure statistical heuristics instead of model-based perplexity scoring.
//!
//! Enhanced with token-aware semantic preservation:
//! - Protects code blocks, JSON, paths, identifiers
//! - Contextual stopword filtering
//! - Preserves negations, comparators, domain terms

use crate::compressor::{CompressionResult, OutputFormat};
#[cfg(feature = "image")]
use crate::image_renderer::{ImageRenderer, ImageRendererConfig};
use regex::Regex;
use std::collections::HashMap;
use std::sync::OnceLock;

/// Type of protected span that should not be modified
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SpanType {
    CodeBlock,    // ```...```
    JsonBlock,    // {...} with high density of :
    Path,         // /path/to/file.ext, http://...
    Identifier,   // camelCase, snake_case, UPPER_SNAKE
    HashOrNumber, // 0x1a2b3c, UUID, large numbers
    Bracket,      // Content inside brackets/braces/parens
}

/// A span of text that should be protected from modification
#[derive(Debug, Clone)]
struct ProtectedSpan {
    start: usize,
    end: usize,
    _span_type: SpanType,
}

/// Importance score for a word based on statistical features
#[derive(Debug, Clone)]
pub struct WordImportance {
    /// Position in the text
    pub position: usize,
    /// Text representation
    pub text: String,
    /// Combined importance score
    pub score: f64,
}

/// Configuration for statistical filtering
#[derive(Debug, Clone)]
pub struct StatisticalFilterConfig {
    /// Target compression ratio (0.0 to 1.0)
    /// 0.5 = keep 50% of tokens, 0.2 = keep 20%
    pub compression_ratio: f32,

    /// Weight for inverse document frequency (IDF)
    pub idf_weight: f32,

    /// Weight for position in document (start/end more important)
    pub position_weight: f32,

    /// Weight for part-of-speech heuristics
    pub pos_weight: f32,

    /// Weight for named entity patterns
    pub entity_weight: f32,

    /// Weight for local entropy (vocabulary diversity)
    pub entropy_weight: f32,

    // Token-aware semantic preservation options
    /// Enable protection masks for code/JSON/paths/identifiers
    pub enable_protection_masks: bool,

    /// Enable contextual stopword filtering (smarter removal)
    pub enable_contextual_stopwords: bool,

    /// Preserve negations (not, no, never, don't, etc.)
    pub preserve_negations: bool,

    /// Preserve comparators (!=, <=, >=, ==, etc.)
    pub preserve_comparators: bool,

    /// Domain-specific terms to always preserve
    pub domain_terms: Vec<String>,

    /// Minimum gap between critical tokens before re-adding
    pub min_gap_between_critical: usize,
}

impl Default for StatisticalFilterConfig {
    fn default() -> Self {
        // Recommended default: 50% compression with 89% quality retention
        // Validated on 20 real papers: 92% keyword retention, 90% entity retention
        // Speed: <0.2ms average
        // Token-aware enhancements: Protects code, contextual stopwords, preserves semantics
        Self {
            compression_ratio: 0.5, // Keep 50% of tokens (recommended)
            idf_weight: 0.3,
            position_weight: 0.2,
            pos_weight: 0.2,
            entity_weight: 0.2,
            entropy_weight: 0.1,
            // Token-aware semantic preservation (all enabled by default)
            enable_protection_masks: true,
            enable_contextual_stopwords: true,
            preserve_negations: true,
            preserve_comparators: true,
            domain_terms: vec![
                "Vectorizer".to_string(),
                "Synap".to_string(),
                "UMICP".to_string(),
                "Graphs".to_string(),
            ],
            min_gap_between_critical: 3,
        }
    }
}

/// Statistical token filter (model-free alternative to LLMLingua)
#[derive(Debug)]
pub struct StatisticalFilter {
    config: StatisticalFilterConfig,
}

impl StatisticalFilter {
    /// Create a new statistical filter
    pub fn new(config: StatisticalFilterConfig) -> Self {
        Self { config }
    }
}

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

impl StatisticalFilter {
    /// Detect JSON spans with proper nesting support (handles multiline and nested JSON)
    fn detect_json_spans(text: &str) -> Vec<(usize, usize)> {
        let mut spans = Vec::new();
        let chars: Vec<char> = text.chars().collect();
        let mut i = 0;
        
        while i < chars.len() {
            // Look for opening brace or bracket
            if chars[i] == '{' || chars[i] == '[' {
                let opening = chars[i];
                let closing = if opening == '{' { '}' } else { ']' };
                let start = i;
                let mut depth = 1;
                let mut in_string = false;
                let mut escape_next = false;
                let mut has_colon = false;
                i += 1;
                
                // Find matching closing brace/bracket
                while i < chars.len() && depth > 0 {
                    if escape_next {
                        escape_next = false;
                        i += 1;
                        continue;
                    }
                    
                    match chars[i] {
                        '\\' if in_string => escape_next = true,
                        '"' => in_string = !in_string,
                        ':' if !in_string => has_colon = true,
                        c if c == opening && !in_string => depth += 1,
                        c if c == closing && !in_string => {
                            depth -= 1;
                            if depth == 0 {
                                // Only add if it looks like JSON (has colons for objects)
                                if opening == '[' || has_colon {
                                    spans.push((start, i + 1));
                                }
                            }
                        }
                        _ => {}
                    }
                    i += 1;
                }
            } else {
                i += 1;
            }
        }
        
        spans
    }

    /// Detect protected spans in text that should not be modified
    fn detect_protected_spans(&self, text: &str) -> Vec<ProtectedSpan> {
        if !self.config.enable_protection_masks {
            return Vec::new();
        }

        let mut spans = Vec::new();

        // Code blocks (```...```)
        static CODE_BLOCK_RE: OnceLock<Regex> = OnceLock::new();
        let code_re = CODE_BLOCK_RE.get_or_init(|| Regex::new(r"```[\s\S]*?```").unwrap());
        for mat in code_re.find_iter(text) {
            spans.push(ProtectedSpan {
                start: mat.start(),
                end: mat.end(),
                _span_type: SpanType::CodeBlock,
            });
        }

        // JSON blocks (improved detection: nested and multiline support)
        // Detect JSON objects and arrays with proper nesting
        let json_spans = Self::detect_json_spans(text);
        for (start, end) in json_spans {
            spans.push(ProtectedSpan {
                start,
                end,
                _span_type: SpanType::JsonBlock,
            });
        }

        // Paths and URLs
        static PATH_RE: OnceLock<Regex> = OnceLock::new();
        let path_re = PATH_RE.get_or_init(|| {
            Regex::new(r"(?:[A-Za-z]+:)?//[^\s]+|[/\\][\w/\\.-]+\.[A-Za-z0-9]{1,5}\b").unwrap()
        });
        for mat in path_re.find_iter(text) {
            spans.push(ProtectedSpan {
                start: mat.start(),
                end: mat.end(),
                _span_type: SpanType::Path,
            });
        }

        // CamelCase identifiers
        static CAMEL_RE: OnceLock<Regex> = OnceLock::new();
        let camel_re =
            CAMEL_RE.get_or_init(|| Regex::new(r"\b[A-Z][a-z0-9]+[A-Z][A-Za-z0-9]+\b").unwrap());
        for mat in camel_re.find_iter(text) {
            spans.push(ProtectedSpan {
                start: mat.start(),
                end: mat.end(),
                _span_type: SpanType::Identifier,
            });
        }

        // snake_case identifiers
        static SNAKE_RE: OnceLock<Regex> = OnceLock::new();
        let snake_re = SNAKE_RE.get_or_init(|| Regex::new(r"\b[a-z_][a-z0-9_]{2,}\b").unwrap());
        for mat in snake_re.find_iter(text) {
            if mat.as_str().contains('_') {
                spans.push(ProtectedSpan {
                    start: mat.start(),
                    end: mat.end(),
                    _span_type: SpanType::Identifier,
                });
            }
        }

        // UPPER_SNAKE_CASE identifiers
        static UPPER_SNAKE_RE: OnceLock<Regex> = OnceLock::new();
        let upper_snake_re =
            UPPER_SNAKE_RE.get_or_init(|| Regex::new(r"\b[A-Z][A-Z0-9_]+\b").unwrap());
        for mat in upper_snake_re.find_iter(text) {
            if mat.as_str().len() > 1 {
                spans.push(ProtectedSpan {
                    start: mat.start(),
                    end: mat.end(),
                    _span_type: SpanType::Identifier,
                });
            }
        }

        // Hashes and large numbers
        static HASH_RE: OnceLock<Regex> = OnceLock::new();
        let hash_re = HASH_RE.get_or_init(|| Regex::new(r"\b[0-9a-f]{7,}\b|\b\d{3,}\b").unwrap());
        for mat in hash_re.find_iter(text) {
            spans.push(ProtectedSpan {
                start: mat.start(),
                end: mat.end(),
                _span_type: SpanType::HashOrNumber,
            });
        }

        // Brackets, braces, parens content
        static BRACKET_RE: OnceLock<Regex> = OnceLock::new();
        let bracket_re =
            BRACKET_RE.get_or_init(|| Regex::new(r"[\{\[\(][^\}\]\)]*[\}\]\)]").unwrap());
        for mat in bracket_re.find_iter(text) {
            spans.push(ProtectedSpan {
                start: mat.start(),
                end: mat.end(),
                _span_type: SpanType::Bracket,
            });
        }

        spans
    }

    /// Check if a word/token position overlaps with any protected span
    fn is_word_protected(
        &self,
        word_start: usize,
        word_end: usize,
        protected: &[ProtectedSpan],
    ) -> bool {
        protected.iter().any(|span| {
            // Check for overlap: word overlaps if it starts before span ends AND ends after span starts
            word_start < span.end && word_end > span.start
        })
    }

    /// Check if a stopword should be preserved based on context
    fn should_preserve_stopword(
        &self,
        word: &str,
        context_before: &[&str],
        context_after: &[&str],
    ) -> bool {
        if !self.config.enable_contextual_stopwords {
            return false;
        }

        let word_lower = word.to_lowercase();

        // "to" in infinitive/phrasal verbs: "how to", "steps to", "need to"
        if word_lower == "to" {
            if let Some(&prev) = context_before.last() {
                let prev_lower = prev.to_lowercase();
                if ["how", "steps", "need", "want", "try", "used", "able"]
                    .contains(&prev_lower.as_str())
                {
                    return true;
                }
            }
        }

        // "in/on/at" followed by paths or technical terms
        if ["in", "on", "at"].contains(&word_lower.as_str()) {
            if let Some(&next) = context_after.first() {
                // Check if next word looks like a path component
                if next.contains('/') || next.contains('\\') || next.contains('.') {
                    return true;
                }
                // Check if next word is technical (starts with uppercase or contains _)
                if next.chars().next().is_some_and(|c| c.is_uppercase()) || next.contains('_') {
                    return true;
                }
            }
        }

        // "is/are/was/were" in assertions (follows important term)
        if ["is", "are", "was", "were", "be"].contains(&word_lower.as_str()) {
            if let Some(&prev) = context_before.last() {
                // If previous word is capitalized or technical, keep the verb
                if prev.chars().next().is_some_and(|c| c.is_uppercase())
                    || prev.len() > 6
                    || prev.contains('_')
                {
                    return true;
                }
            }
        }

        // "and/or" between important terms
        if ["and", "or"].contains(&word_lower.as_str()) {
            let prev_important = context_before.last().is_some_and(|&prev| {
                prev.chars().next().is_some_and(|c| c.is_uppercase()) || prev.len() > 6
            });
            let next_important = context_after.first().is_some_and(|&next| {
                next.chars().next().is_some_and(|c| c.is_uppercase()) || next.len() > 6
            });
            if prev_important && next_important {
                return true;
            }
        }

        false
    }

    /// Check if a word is a critical term that must be preserved
    fn is_critical_term(&self, word: &str) -> Option<f64> {
        let word_lower = word.to_lowercase();

        // Domain-specific terms (highest priority - always preserve)
        for domain_term in &self.config.domain_terms {
            if word.eq_ignore_ascii_case(domain_term) {
                return Some(f64::INFINITY);
            }
        }

        // Negations (very high priority)
        if self.config.preserve_negations {
            const NEGATIONS: &[&str] = &[
                "not",
                "no",
                "never",
                "don't",
                "won't",
                "can't",
                "couldn't",
                "wouldn't",
                "shouldn't",
                "mustn't",
                "haven't",
                "hasn't",
                "hadn't",
                "isn't",
                "aren't",
                "wasn't",
                "weren't",
                "neither",
                "nor",
                "none",
            ];
            if NEGATIONS.contains(&word_lower.as_str()) {
                return Some(10.0);
            }
        }

        // Comparators and operators (very high priority)
        if self.config.preserve_comparators {
            const COMPARATORS: &[&str] = &["!=", "!==", "<=", ">=", "<", ">", "==", "===", "!"];
            if COMPARATORS.contains(&word) {
                return Some(10.0);
            }
        }

        // Modal qualifiers (high priority)
        const MODALS: &[&str] = &[
            "only", "except", "must", "should", "may", "might", "at", "least", "most",
        ];
        if MODALS.contains(&word_lower.as_str()) {
            return Some(5.0);
        }

        None
    }

    /// Calculate importance scores for all tokens
    /// Score words in text by importance
    pub fn score_words(&self, text: &str) -> Vec<WordImportance> {
        let words: Vec<&str> = text.split_whitespace().collect();

        if words.is_empty() {
            return Vec::new();
        }

        // Detect protected spans
        let protected_spans = self.detect_protected_spans(text);

        // Build a mapping of word index to character position in original text
        let word_positions: Vec<(usize, usize)> = {
            let mut positions = Vec::new();
            let mut char_idx = 0;
            let text_chars: Vec<char> = text.chars().collect();

            for word in &words {
                // Skip whitespace
                while char_idx < text_chars.len() && text_chars[char_idx].is_whitespace() {
                    char_idx += 1;
                }

                let start = char_idx;
                let word_len = word.chars().count();
                char_idx += word_len;
                let end = char_idx;

                positions.push((start, end));
            }
            positions
        };

        // Calculate various statistical features
        let idf_scores = self.calculate_idf(&words);
        let position_scores = self.calculate_position_importance(&words);
        let pos_scores = self.calculate_pos_importance(&words, &protected_spans, text);
        let entity_scores = self.calculate_entity_importance(&words);
        let entropy_scores = self.calculate_local_entropy(&words);

        // Combine scores for each word
        words
            .iter()
            .enumerate()
            .map(|(idx, word)| {
                // Check if word is critical or protected
                // First check if it's a critical term
                let final_score = if let Some(critical_score) = self.is_critical_term(word) {
                    critical_score
                } else {
                    // Check if word is in a protected span using the character position
                    let (start, end) = word_positions[idx];
                    let is_protected = self.is_word_protected(start, end, &protected_spans);

                    if is_protected {
                        f64::INFINITY // Never remove protected words
                    } else {
                        // Calculate normal combined score
                        let idf = idf_scores.get(*word).copied().unwrap_or(0.0);
                        let pos_score = position_scores[idx];
                        let pos_tag_score = pos_scores[idx];
                        let entity_score = entity_scores[idx];
                        let entropy = entropy_scores[idx];

                        idf * self.config.idf_weight as f64
                            + pos_score * self.config.position_weight as f64
                            + pos_tag_score * self.config.pos_weight as f64
                            + entity_score * self.config.entity_weight as f64
                            + entropy * self.config.entropy_weight as f64
                    }
                };

                WordImportance {
                    position: idx,
                    text: word.to_string(),
                    score: final_score,
                }
            })
            .collect()
    }

    /// Filter text keeping only high-importance words
    pub fn compress(&self, text: &str) -> String {
        let importances = self.score_words(text);

        if importances.is_empty() {
            return text.to_string();
        }

        // Separate protected (infinite score) from regular words
        let protected_indices: Vec<usize> = importances
            .iter()
            .filter(|imp| imp.score.is_infinite())
            .map(|imp| imp.position)
            .collect();

        let mut regular_words: Vec<_> = importances
            .iter()
            .filter(|imp| !imp.score.is_infinite())
            .cloned()
            .collect();

        // Sort regular words by score (descending)
        regular_words.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Calculate how many regular words to keep (compression ratio applies only to regular words)
        let total_regular = regular_words.len();
        let keep_regular_count = if total_regular > 0 {
            let target_total = (importances.len() as f32 * self.config.compression_ratio) as usize;
            // Subtract protected words from target, ensure we keep at least some regular words
            target_total.saturating_sub(protected_indices.len()).max(1).min(total_regular)
        } else {
            0
        };

        // Get indices of regular tokens to keep
        let mut keep_indices: Vec<usize> = regular_words[..keep_regular_count]
            .iter()
            .map(|imp| imp.position)
            .collect();

        // Add all protected indices (always kept)
        keep_indices.extend(&protected_indices);

        // Fill gaps between critical tokens (using regular_words for gap analysis)
        let critical_threshold = 0.8;
        let mut critical_positions: Vec<usize> = regular_words
            .iter()
            .filter(|imp| imp.score > critical_threshold && keep_indices.contains(&imp.position))
            .map(|imp| imp.position)
            .collect();

        // Also include protected positions in critical positions
        critical_positions.extend(&protected_indices);
        critical_positions.sort_unstable();

        // Check for large gaps between critical tokens
        for window in critical_positions.windows(2) {
            // Ensure window[1] > window[0] to avoid overflow
            if window[1] > window[0] {
                let gap_size = window[1] - window[0];
                if gap_size > self.config.min_gap_between_critical {
                    // Find the highest-scored token in the gap that wasn't kept
                    let gap_candidates: Vec<_> = regular_words
                        .iter()
                        .filter(|imp| {
                            imp.position > window[0]
                                && imp.position < window[1]
                                && !keep_indices.contains(&imp.position)
                        })
                        .collect();

                    if let Some(best_gap_token) = gap_candidates.iter().max_by(|a, b| {
                        a.score
                            .partial_cmp(&b.score)
                            .unwrap_or(std::cmp::Ordering::Equal)
                    }) {
                        keep_indices.push(best_gap_token.position);
                    }
                }
            }
        }

        // Sort by original position to maintain order
        keep_indices.sort_unstable();

        // Reconstruct text with kept tokens
        let words: Vec<&str> = text.split_whitespace().collect();
        keep_indices
            .iter()
            .map(|&idx| words[idx])
            .collect::<Vec<_>>()
            .join(" ")
    }

    /// Compress text and optionally render to image.
    ///
    /// This method performs statistical compression and can output the result
    /// as either plain text or as a 1024x1024 PNG image for vision model consumption.
    ///
    /// # Arguments
    ///
    /// * `text` - The input text to compress
    /// * `format` - Output format (Text or Image)
    ///
    /// # Returns
    ///
    /// A `CompressionResult` containing the compressed text and optional image data.
    ///
    /// # Example
    ///
    /// ```ignore
    /// use compression_prompt::{StatisticalFilter, OutputFormat};
    ///
    /// let filter = StatisticalFilter::default();
    /// let result = filter.compress_with_format("long text...", OutputFormat::Image)?;
    ///
    /// if let Some(img_data) = result.image_data {
    ///     std::fs::write("output.png", img_data)?;
    /// }
    /// ```
    pub fn compress_with_format(
        &self,
        text: &str,
        format: OutputFormat,
    ) -> Result<CompressionResult, Box<dyn std::error::Error>> {
        // Perform statistical compression
        let compressed = self.compress(text);

        // Calculate token counts (rough estimation: words / 4)
        let original_tokens = text.split_whitespace().count();
        let compressed_tokens = compressed.split_whitespace().count();
        let compression_ratio = if original_tokens > 0 {
            compressed_tokens as f32 / original_tokens as f32
        } else {
            1.0
        };
        let tokens_removed = original_tokens.saturating_sub(compressed_tokens);

        // Generate image if requested
        let image_data = if format == OutputFormat::Image {
            #[cfg(feature = "image")]
            {
                let renderer = ImageRenderer::new(ImageRendererConfig::default());
                Some(renderer.render_to_png(&compressed)?)
            }
            #[cfg(not(feature = "image"))]
            {
                None // Image feature not enabled
            }
        } else {
            None
        };

        Ok(CompressionResult {
            compressed,
            image_data,
            format,
            original_tokens,
            compressed_tokens,
            compression_ratio,
            tokens_removed,
        })
    }

    /// Calculate IDF scores
    fn calculate_idf<'a>(&self, words: &[&'a str]) -> HashMap<&'a str, f64> {
        let mut freq_map: HashMap<&str, usize> = HashMap::new();
        for word in words {
            *freq_map.entry(word).or_insert(0) += 1;
        }

        let total = words.len() as f64;
        freq_map
            .iter()
            .map(|(word, count)| (*word, (total / *count as f64).ln()))
            .collect()
    }

    /// Calculate position importance (U-shaped: start and end are important)
    fn calculate_position_importance(&self, words: &[&str]) -> Vec<f64> {
        let len = words.len();
        (0..len)
            .map(|idx| {
                let normalized = idx as f64 / len as f64;
                if !(0.1..=0.9).contains(&normalized) {
                    1.0
                } else if !(0.2..=0.8).contains(&normalized) {
                    0.7
                } else {
                    0.3
                }
            })
            .collect()
    }

    /// Calculate POS importance using stop word heuristics (multilingual)
    /// Supports: English, Spanish, Portuguese, French, German, Italian, Russian,
    /// Chinese, Japanese, Arabic (top 10 world languages)
    /// Enhanced with contextual stopword preservation
    fn calculate_pos_importance(
        &self,
        words: &[&str],
        _protected_spans: &[ProtectedSpan],
        _text: &str,
    ) -> Vec<f64> {
        const STOP_WORDS: &[&str] = &[
            // English
            "the",
            "a",
            "an",
            "and",
            "or",
            "but",
            "in",
            "on",
            "at",
            "to",
            "for",
            "of",
            "with",
            "by",
            "from",
            "as",
            "is",
            "was",
            "are",
            "were",
            "be",
            "been",
            "being",
            "have",
            "has",
            "had",
            "do",
            "does",
            "did",
            "will",
            "would",
            "should",
            "could",
            "may",
            "might",
            "must",
            "can",
            "shall",
            "this",
            "that",
            "these",
            "those",
            "i",
            "you",
            "he",
            "she",
            "it",
            "we",
            "they",
            "what",
            "which",
            "who",
            "when",
            "where",
            "why",
            "how",
            // Spanish (Español)
            "el",
            "la",
            "los",
            "las",
            "un",
            "una",
            "unos",
            "unas",
            "y",
            "o",
            "pero",
            "en",
            "de",
            "del",
            "al",
            "para",
            "por",
            "con",
            "sin",
            "sobre",
            "entre",
            "hasta",
            "desde",
            "es",
            "son",
            "está",
            "están",
            "ser",
            "estar",
            "haber",
            "hacer",
            "tener",
            "decir",
            "ir",
            "ver",
            "dar",
            "saber",
            "querer",
            "poder",
            "poner",
            "este",
            "ese",
            "aquel",
            "mi",
            "tu",
            "su",
            "nuestro",
            "vuestro",
            "que",
            "quien",
            "cual",
            "cuando",
            "donde",
            "como",
            // Portuguese (Português)
            "o",
            "a",
            "os",
            "as",
            "um",
            "uma",
            "uns",
            "umas",
            "e",
            "ou",
            "mas",
            "em",
            "de",
            "do",
            "da",
            "dos",
            "das",
            "no",
            "na",
            "nos",
            "nas",
            "ao",
            "à",
            "aos",
            "às",
            "para",
            "por",
            "com",
            "sem",
            "sobre",
            "entre",
            "até",
            "desde",
            "é",
            "são",
            "está",
            "estão",
            "ser",
            "estar",
            "haver",
            "ter",
            "fazer",
            "dizer",
            "ir",
            "ver",
            "dar",
            "saber",
            "querer",
            "poder",
            "pôr",
            "este",
            "esse",
            "aquele",
            "meu",
            "teu",
            "seu",
            "nosso",
            "vosso",
            "que",
            "quem",
            "qual",
            "quando",
            "onde",
            "como",
            // French (Français)
            "le",
            "la",
            "les",
            "un",
            "une",
            "des",
            "et",
            "ou",
            "mais",
            "dans",
            "en",
            "de",
            "du",
            "au",
            "aux",
            "pour",
            "par",
            "avec",
            "sans",
            "sur",
            "sous",
            "entre",
            "vers",
            "chez",
            "est",
            "sont",
            "être",
            "avoir",
            "faire",
            "dire",
            "aller",
            "voir",
            "savoir",
            "pouvoir",
            "vouloir",
            "venir",
            "devoir",
            "prendre",
            "ce",
            "cet",
            "cette",
            "ces",
            "mon",
            "ton",
            "son",
            "notre",
            "votre",
            "leur",
            "que",
            "qui",
            "quoi",
            "dont",
            "où",
            "quand",
            "comment",
            // German (Deutsch)
            "der",
            "die",
            "das",
            "den",
            "dem",
            "des",
            "ein",
            "eine",
            "einer",
            "eines",
            "einem",
            "einen",
            "und",
            "oder",
            "aber",
            "in",
            "im",
            "an",
            "auf",
            "für",
            "von",
            "zu",
            "mit",
            "bei",
            "nach",
            "über",
            "unter",
            "ist",
            "sind",
            "war",
            "waren",
            "sein",
            "haben",
            "werden",
            "können",
            "müssen",
            "sollen",
            "wollen",
            "dieser",
            "jener",
            "mein",
            "dein",
            "sein",
            "unser",
            "euer",
            "ihr",
            "was",
            "wer",
            "wo",
            "wann",
            "wie",
            "warum",
            // Italian (Italiano)
            "il",
            "lo",
            "l",
            "i",
            "gli",
            "la",
            "le",
            "un",
            "uno",
            "una",
            "e",
            "o",
            "ma",
            "in",
            "di",
            "del",
            "dello",
            "della",
            "dei",
            "degli",
            "delle",
            "al",
            "allo",
            "alla",
            "ai",
            "agli",
            "alle",
            "per",
            "da",
            "dal",
            "dallo",
            "dalla",
            "dai",
            "dagli",
            "dalle",
            "con",
            "su",
            "sul",
            "sullo",
            "sulla",
            "sui",
            "sugli",
            "sulle",
            "è",
            "sono",
            "essere",
            "avere",
            "fare",
            "dire",
            "andare",
            "vedere",
            "sapere",
            "potere",
            "volere",
            "questo",
            "quello",
            "mio",
            "tuo",
            "suo",
            "nostro",
            "vostro",
            "loro",
            "che",
            "chi",
            "quale",
            "quando",
            "dove",
            "come",
            "perché",
            // Russian (Русский) - romanized
            "i",
            "v",
            "ne",
            "na",
            "ya",
            "on",
            "s",
            "eto",
            "kak",
            "po",
            "no",
            "oni",
            "vse",
            "tak",
            "ego",
            "za",
            "byl",
            "bylo",
            "tem",
            "chto",
            "eto",
            "esli",
            "mogu",
            "mozhet",
            "by",
            // Chinese (中文) - common particles
            "的",
            "了",
            "和",
            "是",
            "在",
            "我",
            "有",
            "他",
            "这",
            "中",
            "大",
            "来",
            "上",
            "国",
            "个",
            "到",
            "说",
            "们",
            "为",
            "子",
            "中",
            "你",
            "地",
            "出",
            "道",
            "也",
            "时",
            "年",
            // Japanese (日本語) - particles and common words
            "は",
            "が",
            "を",
            "に",
            "で",
            "と",
            "の",
            "も",
            "や",
            "から",
            "まで",
            "より",
            "か",
            "な",
            "ね",
            "よ",
            "わ",
            "さ",
            "だ",
            "です",
            "ます",
            "ある",
            "いる",
            "する",
            "なる",
            "これ",
            "それ",
            "あれ",
            "この",
            "その",
            "あの",
            "ここ",
            "そこ",
            "あそこ",
            // Arabic (العربية) - romanized common words
            "al",
            "wa",
            "fi",
            "min",
            "ila",
            "an",
            "ma",
            "la",
            "li",
            "bi",
            "qad",
            "lam",
            "kan",
            "fi",
            "ala",
            "hatha",
            "dhalika",
            "huwa",
            "hiya",
            "hum",
            // Hindi (हिन्दी) - romanized common words
            "ka",
            "ki",
            "ke",
            "se",
            "ne",
            "ko",
            "me",
            "par",
            "hai",
            "tha",
            "the",
            "thi",
            "aur",
            "ya",
            "to",
            "is",
            "wo",
            "ye",
            "kya",
            "kaise",
            "kab",
            "kahan",
            "kyun",
        ];

        words
            .iter()
            .enumerate()
            .map(|(idx, word)| {
                let lower = word.to_lowercase();

                // Check if it's a stopword
                if STOP_WORDS.contains(&lower.as_str()) {
                    // Check contextual preservation
                    let context_before: Vec<&str> = if idx > 0 {
                        words[..idx].iter().rev().take(3).rev().copied().collect()
                    } else {
                        Vec::new()
                    };

                    let context_after: Vec<&str> = if idx + 1 < words.len() {
                        words[idx + 1..].iter().take(3).copied().collect()
                    } else {
                        Vec::new()
                    };

                    if self.should_preserve_stopword(word, &context_before, &context_after) {
                        0.7 // Contextually important stopword
                    } else {
                        0.1 // Regular stopword - low importance
                    }
                } else if word.chars().next().is_some_and(|c| c.is_uppercase()) {
                    1.0 // Proper noun - high importance
                } else if word.len() > 6 {
                    0.7 // Long word - medium-high importance
                } else {
                    0.5 // Regular word - medium importance
                }
            })
            .collect()
    }

    /// Detect named entities using simple patterns
    fn calculate_entity_importance(&self, words: &[&str]) -> Vec<f64> {
        words
            .iter()
            .enumerate()
            .map(|(idx, word)| {
                let mut score: f64 = 0.0;
                if word.chars().next().is_some_and(|c| c.is_uppercase()) {
                    score += 0.3;
                }
                if idx > 0 {
                    let prev = words[idx - 1].to_lowercase();
                    if prev.starts_with("mr.") || prev.starts_with("dr.") {
                        score += 0.5;
                    }
                }
                if word.contains('@') || word.starts_with("http") {
                    score += 0.6;
                }
                if word.len() > 1 && word.chars().all(|c| c.is_uppercase()) {
                    score += 0.4;
                }
                score.min(1.0)
            })
            .collect()
    }

    /// Calculate local entropy (vocabulary diversity)
    fn calculate_local_entropy(&self, words: &[&str]) -> Vec<f64> {
        const WINDOW: usize = 10;
        (0..words.len())
            .map(|idx| {
                let start = idx.saturating_sub(WINDOW / 2);
                let end = (idx + WINDOW / 2).min(words.len());
                let window = &words[start..end];
                let unique: std::collections::HashSet<_> = window.iter().collect();
                unique.len() as f64 / window.len() as f64
            })
            .collect()
    }
}

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

    #[test]
    fn test_compression() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.5,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "The quick brown fox jumps over the lazy dog";
        let compressed = filter.compress(text);

        let original_words = text.split_whitespace().count();
        let compressed_words = compressed.split_whitespace().count();

        assert!(compressed_words <= original_words);
        assert!(!compressed.is_empty());
    }

    #[test]
    fn test_code_block_protection() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.3,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "Here is some code ```rust fn main() { println!(\"Hello\"); }``` that should be preserved";
        let compressed = filter.compress(text);

        // Code block should be in the output even with aggressive compression
        assert!(
            compressed.contains("```rust") || compressed.contains("println!"),
            "Expected code block to be preserved, got: {}",
            compressed
        );
    }

    #[test]
    fn test_json_protection() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.3,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "The config is {\"key\": \"value\"} and it should remain intact";
        let compressed = filter.compress(text);

        // JSON should be preserved completely
        assert!(
            compressed.contains("{\"key\":") || compressed.contains("\"key\""),
            "Expected JSON to be preserved, got: {}",
            compressed
        );
    }

    #[test]
    fn test_nested_json_protection() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.2,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "Here is a nested JSON {\"user\": {\"name\": \"John\", \"age\": 30}, \"active\": true} that must be kept";
        let compressed = filter.compress(text);

        // Nested JSON should be preserved
        assert!(
            compressed.contains("{\"user\":") || compressed.contains("\"name\"") || compressed.contains("John"),
            "Expected nested JSON to be preserved, got: {}",
            compressed
        );
    }

    #[test]
    fn test_multiline_json_protection() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.2,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = r#"Configuration:
{
  "host": "localhost",
  "port": 8080,
  "options": {
    "debug": true
  }
}
End of config"#;
        let compressed = filter.compress(text);

        // Multiline JSON should be preserved
        assert!(
            compressed.contains("\"host\"") || compressed.contains("localhost") || compressed.contains("8080"),
            "Expected multiline JSON to be preserved, got: {}",
            compressed
        );
    }

    #[test]
    fn test_json_array_protection() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.2,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "The list is [1, 2, 3, 4, 5] and should be preserved";
        let compressed = filter.compress(text);

        // JSON array should be preserved
        assert!(
            compressed.contains("[1,") || compressed.contains("2") || compressed.contains("5"),
            "Expected JSON array to be preserved, got: {}",
            compressed
        );
    }

    #[test]
    fn test_path_preservation() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.4,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "Check the file in src/main.rs for the implementation details";
        let compressed = filter.compress(text);

        // Path should be preserved
        assert!(
            compressed.contains("src/main.rs")
                || (compressed.contains("src") && compressed.contains("main.rs"))
        );
    }

    #[test]
    fn test_contextual_stopword_to() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.5,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        // "to" should be kept in "how to"
        let text1 = "how to reproduce the bug";
        let compressed1 = filter.compress(text1);
        assert!(compressed1.contains("to") || compressed1.contains("how"));

        // "to" can be removed in other contexts if not critical
        let text2 = "going to the store";
        let _compressed2 = filter.compress(text2);
        // This is context-dependent, so we don't assert removal
    }

    #[test]
    fn test_negation_preservation() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.3,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "do not remove this critical information";
        let compressed = filter.compress(text);

        // "not" should always be preserved
        assert!(compressed.contains("not"));
    }

    #[test]
    fn test_comparator_preservation() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.3,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "check if x >= 5 before proceeding";
        let compressed = filter.compress(text);

        // ">=" should be preserved
        assert!(compressed.contains(">=") || compressed.contains("5") || compressed.contains("x"));
    }

    #[test]
    fn test_domain_terms_preservation() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.3,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "use the Vectorizer tool to process data";
        let compressed = filter.compress(text);

        // Domain term "Vectorizer" should be preserved
        assert!(compressed.contains("Vectorizer"));
    }

    #[test]
    fn test_identifier_protection() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.3,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "call the getUserData function from user_service module";
        let compressed = filter.compress(text);

        // Identifiers should be preserved
        assert!(compressed.contains("getUserData") || compressed.contains("user_service"));
    }

    #[test]
    fn test_gap_filling_between_critical_tokens() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.2,
            min_gap_between_critical: 2,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "Vectorizer is a critical component that handles data processing for Synap";
        let compressed = filter.compress(text);

        // Should have some words between Vectorizer and Synap
        assert!(
            compressed.contains("Vectorizer"),
            "Expected 'Vectorizer' in output: {}",
            compressed
        );
        assert!(
            compressed.contains("Synap"),
            "Expected 'Synap' in output: {}",
            compressed
        );

        let words: Vec<&str> = compressed.split_whitespace().collect();
        assert!(
            words.len() >= 3,
            "Expected at least 3 words, got: {}",
            words.len()
        );
    }

    #[test]
    fn test_protection_masks_can_be_disabled() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.3,
            enable_protection_masks: false,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "Check src/main.rs for details";
        let _compressed = filter.compress(text);

        // With protection disabled, behavior is normal compression
        // Just ensure it doesn't crash
    }

    #[test]
    fn test_contextual_stopwords_can_be_disabled() {
        let config = StatisticalFilterConfig {
            compression_ratio: 0.5,
            enable_contextual_stopwords: false,
            ..Default::default()
        };
        let filter = StatisticalFilter::new(config);

        let text = "how to reproduce the issue";
        let _compressed = filter.compress(text);

        // With contextual stopwords disabled, "to" might be removed
        // Just ensure it doesn't crash
    }
}