//! Lente de Corte Causal (LCC) — atención dispersa subcuadrática nativa.
#![allow(dead_code)]

use std::collections::{HashMap, HashSet};

use bimoh::Memcon;
use lnhash::{text_to_togens, Togen};

/// Mapa de posiciones de tokens en el texto original
pub struct TokenPositionMap {
    pub positions: HashMap<usize, (usize, usize)>,
    pub text: String,
}

/// Resultado de la selección dispersa
pub struct SparseOutput {
    pub compact_text: String,
    pub original_indices: Vec<usize>,
    pub extraction_ratio: f64,
    pub tokens_selected: usize,
    pub tokens_total: usize,
}

#[derive(Debug, Clone)]
struct EnhancedMemcon {
    memcon: Memcon,
    relevance_score: f64,
    token_ranges: Vec<(usize, usize)>,
}

pub fn lcc_slice(
    full_text: &str,
    memcons: &[Memcon],
    token_positions: &TokenPositionMap,
    target_sparsity: f64,
) -> SparseOutput {
    let togens = text_to_togens(full_text);
    let token_count = togens.len();
    if token_count == 0 || memcons.is_empty() {
        return SparseOutput {
            compact_text: String::new(),
            original_indices: Vec::new(),
            extraction_ratio: 0.0,
            tokens_selected: 0,
            tokens_total: 0,
        };
    }

    let enhanced = associate_memcons_with_ranges(memcons, &togens, token_positions);

    let mut sorted = enhanced;
    sorted.sort_by(|a, b| b.relevance_score.partial_cmp(&a.relevance_score).unwrap());

    let target_tokens = ((token_count as f64) * target_sparsity) as usize;
    let mut selected = Vec::new();
    let mut acc = 0;

    for memcon in sorted {
        let count = memcon.token_ranges.len();
        if acc + count <= target_tokens {
            selected.push(memcon);
            acc += count;
        } else {
            let needed = target_tokens - acc;
            if let Some(partial) = create_partial(&memcon, needed) {
                selected.push(partial);
                acc += needed;
            }
            break;
        }
        if acc >= target_tokens { break; }
    }

    let (compact_text, original_indices) = extract_fragments(&selected, &token_positions.text);
    let ratio = if token_count > 0 { acc as f64 / token_count as f64 } else { 0.0 };

    SparseOutput {
        compact_text,
        original_indices,
        extraction_ratio: ratio,
        tokens_selected: acc,
        tokens_total: token_count,
    }
}

const CONTEXT_WINDOW: usize = 2;
fn associate_memcons_with_ranges(
    memcons: &[Memcon],
    togens: &[Togen],
    token_positions: &TokenPositionMap,
) -> Vec<EnhancedMemcon> {
    if memcons.is_empty() { return Vec::new(); }

    let memcon_ids: HashSet<u32> = memcons.iter().map(|m| m.id).collect();
    let mut id_to_indices: HashMap<u32, Vec<usize>> = HashMap::new();
    for (idx, togen) in togens.iter().enumerate() {
        if memcon_ids.contains(&togen.id) {
            id_to_indices.entry(togen.id).or_default().push(idx);
        }
    }

    if id_to_indices.is_empty() {
        let chunk = togens.len() / memcons.len().max(1);
        return memcons.iter().enumerate().map(|(i, m)| {
            let start = i * chunk;
            let end = ((i + 1) * chunk).min(togens.len());
            let ranges = if start < end { vec![(start, end - 1)] } else { vec![] };
            EnhancedMemcon {
                memcon: m.clone(),
                relevance_score: m.complexity as f64,
                token_ranges: ranges,
            }
        }).collect();
    }

    let total_tokens = token_positions.text.split_whitespace().count().max(1);

    memcons.iter().map(|memcon| {
        let indices = id_to_indices.get(&memcon.id).cloned().unwrap_or_default();
        if indices.is_empty() {
            return EnhancedMemcon {
                memcon: memcon.clone(),
                relevance_score: 0.0,
                token_ranges: vec![],
            };
        }

        let mut expanded: Vec<(usize, usize)> = indices.iter()
            .map(|&idx| {
                let s = idx.saturating_sub(CONTEXT_WINDOW);
                let e = (idx + CONTEXT_WINDOW).min(total_tokens - 1);
                (s, e)
            })
            .collect();
        expanded.sort_unstable();

        let mut merged = Vec::new();
        let mut cur = expanded[0];
        for &(s, e) in &expanded[1..] {
            if s <= cur.1 + 1 { cur.1 = cur.1.max(e); }
            else { merged.push(cur); cur = (s, e); }
        }
        merged.push(cur);

        EnhancedMemcon {
            memcon: memcon.clone(),
            relevance_score: memcon.complexity as f64,
            token_ranges: merged,
        }
    }).collect()
}

fn create_partial(memcon: &EnhancedMemcon, needed: usize) -> Option<EnhancedMemcon> {
    if needed == 0 { return None; }
    let n = needed.min(memcon.token_ranges.len());
    Some(EnhancedMemcon {
        memcon: memcon.memcon.clone(),
        relevance_score: memcon.relevance_score,
        token_ranges: memcon.token_ranges[..n].to_vec(),
    })
}

fn extract_fragments(memcons: &[EnhancedMemcon], full_text: &str) -> (String, Vec<usize>) {
    let mut text = String::new();
    let mut indices = Vec::new();

    let mut sorted = memcons.to_vec();
    sorted.sort_by(|a, b| {
        a.token_ranges.first().unwrap_or(&(0, 0)).0
            .cmp(&b.token_ranges.first().unwrap_or(&(0, 0)).0)
    });

    let words: Vec<&str> = full_text.split_whitespace().collect();
    for memcon in sorted {
        for &(start, end) in &memcon.token_ranges {
            let s = start.min(words.len().saturating_sub(1));
            let e = end.min(words.len().saturating_sub(1));
            if s <= e {
                if !text.is_empty() { text.push(' '); }
                text.push_str(&words[s..=e].join(" "));
                indices.push(s);
            }
        }
    }
    (text, indices)
}

pub fn build_position_map(text: &str) -> TokenPositionMap {
    let words: Vec<&str> = text.split_whitespace().collect();
    let mut positions = HashMap::new();
    let mut ci = 0;
    for (ti, w) in words.iter().enumerate() {
        positions.insert(ti, (ci, ci + w.len()));
        ci += w.len() + 1;
    }
    TokenPositionMap { positions, text: text.to_string() }
}

/// Sparsity adaptativa basada en cantidad de tokens y complejidad de memcons.
///
/// La fórmula ancla la compresión en la riqueza semántica del texto:
/// - Muchos memcons con alta complejidad (d_i > 2.0) → más conservador (conservas más)
/// - Pocos memcons o baja complejidad → compresión más agresiva
///
/// S = max(S_min, min(1.0, K / (tokens * C_avg)))
/// donde C_avg = promedio de complexity de los memcons, K ≈ 50 (calibrable)
///
/// S_min escala con tokens: 1% para 2000 tokens, 0.5% para 100k, floor 0.3%
pub fn dynamic_sparsity(token_count: usize) -> f64 {
    if token_count < 2000 { return 1.0; }
    (50.0 / token_count as f64).max(0.003)
}

/// Sparsity adaptativa que considera la complejidad de los memcons consolidados.
/// A mayor complejidad promedio, mayor retención (sparsidad más alta).
pub fn adaptive_sparsity(token_count: usize, complexities: &[f32]) -> f64 {
    if token_count < 2000 || complexities.is_empty() {
        return 1.0;
    }

    let avg_complexity: f64 = complexities.iter()
        .map(|&c| c as f64)
        .sum::<f64>() / complexities.len() as f64;

    let memcon_factor = (complexities.len() as f64).min(64.0);
    let complexity_factor = avg_complexity.max(0.3);

    // K base calibrado: 50 para textos típicos
    // Ajustado por: más memcons → más tokens necesarios → sparsidad más alta
    //               mayor complejidad → conceptos más densos → necesitas más contexto
    let k = 50.0 * (1.0 + memcon_factor / 32.0) * complexity_factor;

    let raw_sparsity = k / token_count as f64;

    // Floor dinámico: 0.3% mínimo absoluto, escala con complejidad
    let floor = if avg_complexity > 2.0 {
        0.008  // 0.8% para textos muy densos
    } else if avg_complexity > 1.5 {
        0.005  // 0.5%
    } else {
        0.003  // 0.3%
    };

    raw_sparsity.max(floor).min(1.0)
}

/// Contar tokens (wrapper simple)
pub fn count_tokens(text: &str) -> usize {
    text.split_whitespace().count()
}

/// Convierte IDs hex + complejidades a Vec<Memcon>
pub fn memcons_from_hex(hex_ids: &[String], complexities: &[f32]) -> Vec<Memcon> {
    hex_ids.iter()
        .zip(complexities.iter().copied().chain(std::iter::repeat(0.5)))
        .filter_map(|(h, cplx)| u32::from_str_radix(h, 16).ok().map(|id| Memcon { id, complexity: cplx }))
        .collect()
}

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

    #[test]
    fn test_lcc_slice() {
        let text = "Este es un texto de prueba para el selector disperso. Contiene varias oraciones.";
        let memcons = vec![
            Memcon { id: 1, complexity: 0.5 },
            Memcon { id: 2, complexity: 0.8 },
        ];
        let pm = build_position_map(text);
        let r = lcc_slice(text, &memcons, &pm, 0.13);
        assert!(r.extraction_ratio > 0.0);
        assert!(r.tokens_selected <= r.tokens_total);
    }

    #[test]
    fn test_empty() {
        let text = "Hello test";
        let pm = build_position_map(text);
        let r = lcc_slice(text, &[], &pm, 0.13);
        assert_eq!(r.tokens_selected, 0);
    }

    #[test]
    fn test_position_map() {
        let text = "Hola mundo";
        let pm = build_position_map(text);
        assert_eq!(pm.positions.len(), 2);
    }

    #[test]
    fn test_sparsity() {
        assert!((dynamic_sparsity(500) - 1.0).abs() < 0.001, "<2000 returns 1.0");
        assert!((dynamic_sparsity(3000) - 0.0167).abs() < 0.01, "3000 -> 50/3000 ~ 0.0167");
        assert!((dynamic_sparsity(100000) - 0.003).abs() < 0.001, "large -> floor 0.3%");
    }

    #[test]
    fn test_adaptive_sparsity_short_text() {
        assert!((adaptive_sparsity(500, &[0.5, 0.8]) - 1.0).abs() < 0.001);
    }

    #[test]
    fn test_adaptive_sparsity_high_complexity() {
        // 5000 tokens, 3 memcons con alta complejidad
        let s = adaptive_sparsity(5000, &[2.5, 2.8, 3.0]);
        assert!(s > 0.01, "high complexity → more retention: got {s}");
        assert!(s <= 1.0);
    }

    #[test]
    fn test_adaptive_sparsity_low_complexity() {
        // 5000 tokens, 2 memcons débiles
        let s = adaptive_sparsity(5000, &[0.4, 0.5]);
        assert!(s < 0.01, "low complexity → aggressive compression: got {s}");
        assert!(s >= 0.003, "never below floor 0.3%: got {s}");
    }

    #[test]
    fn test_adaptive_sparsity_empty_complexities() {
        assert!((adaptive_sparsity(5000, &[]) - 1.0).abs() < 0.001);
    }

    #[test]
    fn test_adaptive_sparsity_many_memcons() {
        // 10k tokens, 40 memcons de alta complejidad → debería ser más conservador
        let many: Vec<f32> = (0..40).map(|_| 2.5).collect();
        let few: Vec<f32> = vec![2.5, 2.5];
        let s_many = adaptive_sparsity(10000, &many);
        let s_few = adaptive_sparsity(10000, &few);
        assert!(s_many >= s_few, "more memcons → equal or more retention: {s_many} vs {s_few}");
    }
}
