remem-ai 0.5.145

use std::collections::{BTreeMap, BTreeSet};

use super::types::{MemoryBenchCondition, MemoryBenchEvidence, MemoryBenchTask};

const TOP_K: usize = 5;
const TRUNCATED_CONTEXT_ITEMS: usize = 1;

pub(super) fn fixture_retrieval_indices(
    condition: MemoryBenchCondition,
    task: &MemoryBenchTask,
) -> Option<Vec<usize>> {
    let indices = match condition {
        MemoryBenchCondition::NoMemory => Vec::new(),
        MemoryBenchCondition::TruncatedFullContext => recent_allowed(task, TRUNCATED_CONTEXT_ITEMS),
        MemoryBenchCondition::OracleEvidence => gold_supporting(task),
        MemoryBenchCondition::CompleteStoredMemory => allowed_evidence(task),
        MemoryBenchCondition::Bm25Baseline => rank_by_bm25(task),
        MemoryBenchCondition::VectorBaseline => rank_by_vector_proxy(task),
        MemoryBenchCondition::HybridRagBaseline => rank_by_hybrid_rag(task),
        MemoryBenchCondition::SummaryBaseline => active_summary_evidence(task),
        MemoryBenchCondition::RetrievedMemory | MemoryBenchCondition::RememDefault => {
            return None;
        }
    };
    Some(indices)
}

fn gold_supporting(task: &MemoryBenchTask) -> Vec<usize> {
    task.evidence
        .iter()
        .enumerate()
        .filter(|(_, evidence)| task.gold_supporting_event_ids.contains(&evidence.event_id))
        .map(|(idx, _)| idx)
        .collect()
}

fn allowed_evidence(task: &MemoryBenchTask) -> Vec<usize> {
    task.evidence
        .iter()
        .enumerate()
        .filter(|(_, evidence)| evidence.retention_allowed)
        .map(|(idx, _)| idx)
        .collect()
}

fn recent_allowed(task: &MemoryBenchTask, limit: usize) -> Vec<usize> {
    let mut scored = task
        .evidence
        .iter()
        .enumerate()
        .filter(|(_, evidence)| evidence.retention_allowed)
        .map(|(idx, evidence)| (idx, evidence.created_at_epoch.unwrap_or_default()))
        .collect::<Vec<_>>();
    scored.sort_by(|left, right| right.1.cmp(&left.1).then_with(|| left.0.cmp(&right.0)));
    scored.into_iter().take(limit).map(|(idx, _)| idx).collect()
}

fn active_summary_evidence(task: &MemoryBenchTask) -> Vec<usize> {
    task.evidence
        .iter()
        .enumerate()
        .filter(|(_, evidence)| evidence.retention_allowed && evidence.status == "active")
        .map(|(idx, _)| idx)
        .take(TOP_K)
        .collect()
}

fn rank_by_bm25(task: &MemoryBenchTask) -> Vec<usize> {
    rank_allowed(task, |evidence| lexical_score(&task.query, evidence))
}

fn rank_by_vector_proxy(task: &MemoryBenchTask) -> Vec<usize> {
    rank_allowed(task, |evidence| vector_proxy_score(&task.query, evidence))
}

fn rank_by_hybrid_rag(task: &MemoryBenchTask) -> Vec<usize> {
    let bm25 = ranked_scores(task, |evidence| lexical_score(&task.query, evidence));
    let vector = ranked_scores(task, |evidence| vector_proxy_score(&task.query, evidence));
    let mut fused: BTreeMap<usize, f64> = BTreeMap::new();
    for (rank, (idx, _)) in bm25.iter().enumerate() {
        *fused.entry(*idx).or_default() += rrf_score(rank);
    }
    for (rank, (idx, _)) in vector.iter().enumerate() {
        *fused.entry(*idx).or_default() += rrf_score(rank);
    }
    let mut ranked = fused.into_iter().collect::<Vec<_>>();
    ranked.sort_by(|left, right| {
        right
            .1
            .total_cmp(&left.1)
            .then_with(|| left.0.cmp(&right.0))
    });
    ranked.into_iter().take(TOP_K).map(|(idx, _)| idx).collect()
}

fn rrf_score(rank: usize) -> f64 {
    1.0 / (60.0 + rank as f64 + 1.0)
}

fn rank_allowed(task: &MemoryBenchTask, score: impl Fn(&MemoryBenchEvidence) -> f64) -> Vec<usize> {
    ranked_scores(task, score)
        .into_iter()
        .take(TOP_K)
        .map(|(idx, _)| idx)
        .collect()
}

fn ranked_scores(
    task: &MemoryBenchTask,
    score: impl Fn(&MemoryBenchEvidence) -> f64,
) -> Vec<(usize, f64)> {
    let mut scored = task
        .evidence
        .iter()
        .enumerate()
        .filter(|(_, evidence)| evidence.retention_allowed)
        .map(|(idx, evidence)| {
            let status_bonus = if evidence.status == "active" {
                0.25
            } else {
                0.0
            };
            (idx, score(evidence) + status_bonus)
        })
        .collect::<Vec<_>>();
    scored.sort_by(|left, right| {
        right
            .1
            .total_cmp(&left.1)
            .then_with(|| left.0.cmp(&right.0))
    });
    scored
}

fn lexical_score(query: &str, evidence: &MemoryBenchEvidence) -> f64 {
    let query_tokens = tokens(query);
    if query_tokens.is_empty() {
        return 0.0;
    }
    let title_tokens = tokens(&evidence.title);
    let content_tokens = tokens(&evidence.content);
    let title_hits = query_tokens
        .iter()
        .filter(|token| title_tokens.contains(*token))
        .count();
    let content_hits = query_tokens
        .iter()
        .filter(|token| content_tokens.contains(*token))
        .count();
    (title_hits as f64 * 2.0) + content_hits as f64
}

fn vector_proxy_score(query: &str, evidence: &MemoryBenchEvidence) -> f64 {
    let query_grams = char_trigrams(query);
    if query_grams.is_empty() {
        return 0.0;
    }
    let mut document_grams = char_trigrams(&evidence.title);
    document_grams.extend(char_trigrams(&evidence.content));
    let intersection = query_grams.intersection(&document_grams).count();
    let union = query_grams.union(&document_grams).count();
    if union == 0 {
        0.0
    } else {
        intersection as f64 / union as f64
    }
}

fn tokens(value: &str) -> BTreeSet<String> {
    value
        .split(|ch: char| !ch.is_ascii_alphanumeric())
        .filter_map(|token| {
            let normalized = token.trim().to_ascii_lowercase();
            (normalized.len() >= 2).then_some(normalized)
        })
        .collect()
}

fn char_trigrams(value: &str) -> BTreeSet<String> {
    let chars = value
        .to_ascii_lowercase()
        .chars()
        .filter(|ch| ch.is_ascii_alphanumeric())
        .collect::<Vec<_>>();
    chars
        .windows(3)
        .map(|window| window.iter().collect::<String>())
        .collect()
}