//! Repetition detection and penalty implementation

use std::collections::HashMap;

/// Advanced repetition detection
pub struct RepetitionDetector {
    ngram_size: usize,
    window_size: usize,
}

impl RepetitionDetector {
    pub fn new(ngram_size: usize, window_size: usize) -> Self {
        Self {
            ngram_size,
            window_size,
        }
    }

    pub fn detect_repetitions(&self, tokens: &[u32]) -> RepetitionReport {
        let mut ngram_counts = HashMap::new();
        let mut repetitive_spans = Vec::new();

        // Count n-grams in sliding windows
        for window_start in 0..tokens.len().saturating_sub(self.window_size) {
            let window_end = (window_start + self.window_size).min(tokens.len());
            let window = &tokens[window_start..window_end];

            for ngram in window.windows(self.ngram_size) {
                *ngram_counts.entry(ngram.to_vec()).or_insert(0) += 1;
            }
        }

        // Find repetitive patterns
        for (ngram, count) in ngram_counts {
            if count > 2 {
                repetitive_spans.push(RepetitiveSpan {
                    ngram,
                    count,
                    severity: self.calculate_severity(count),
                });
            }
        }

        let overall_score = self.calculate_overall_score(&repetitive_spans);
        RepetitionReport {
            total_repetitions: repetitive_spans.len(),
            repetitive_spans,
            overall_score,
        }
    }

    fn calculate_severity(&self, count: usize) -> RepetitionSeverity {
        match count {
            3..=5 => RepetitionSeverity::Mild,
            6..=10 => RepetitionSeverity::Moderate,
            _ => RepetitionSeverity::Severe,
        }
    }

    fn calculate_overall_score(&self, spans: &[RepetitiveSpan]) -> f32 {
        if spans.is_empty() {
            return 1.0;
        }

        let severity_sum: f32 = spans
            .iter()
            .map(|span| match span.severity {
                RepetitionSeverity::Mild => 0.1,
                RepetitionSeverity::Moderate => 0.3,
                RepetitionSeverity::Severe => 0.5,
            })
            .sum();

        (1.0 - severity_sum / spans.len() as f32).clamp(0.0, 1.0)
    }
}

#[derive(Debug, Clone)]
pub struct RepetitionReport {
    pub total_repetitions: usize,
    pub repetitive_spans: Vec<RepetitiveSpan>,
    pub overall_score: f32,
}

#[derive(Debug, Clone)]
pub struct RepetitiveSpan {
    pub ngram: Vec<u32>,
    pub count: usize,
    pub severity: RepetitionSeverity,
}

#[derive(Debug, Clone, PartialEq)]
pub enum RepetitionSeverity {
    Mild,
    Moderate,
    Severe,
}

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

    #[test]
    fn test_repetition_detection() {
        let detector = RepetitionDetector::new(3, 20);
        let tokens = vec![1, 2, 3, 1, 2, 3, 1, 2, 3, 4, 5, 6];

        let report = detector.detect_repetitions(&tokens);
        assert!(report.total_repetitions > 0);
        assert!(report.overall_score < 1.0);
    }
}