kibble 0.1.0

//! Hand-rolled BM25 lexical scorer (no external dependency).

use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};

/// Lowercase, split on non-alphanumeric, drop empties.
pub fn tokenize(text: &str) -> Vec<String> {
    text.split(|c: char| !c.is_alphanumeric())
        .filter(|s| !s.is_empty())
        .map(|s| s.to_lowercase())
        .collect()
}

/// Serializable inverted index + length stats for BM25.
#[derive(Serialize, Deserialize, Default)]
pub struct Bm25Index {
    pub n: usize,
    pub avg_len: f32,
    pub doc_len: Vec<usize>,
    pub postings: HashMap<String, Vec<(usize, u32)>>,
}

impl Bm25Index {
    pub fn build(chunk_tokens: &[Vec<String>]) -> Self {
        let n = chunk_tokens.len();
        let doc_len: Vec<usize> = chunk_tokens.iter().map(|t| t.len()).collect();
        let total: usize = doc_len.iter().sum();
        let avg_len = if n > 0 { total as f32 / n as f32 } else { 0.0 };
        let mut postings: HashMap<String, Vec<(usize, u32)>> = HashMap::new();
        for (idx, toks) in chunk_tokens.iter().enumerate() {
            let mut tf: HashMap<&str, u32> = HashMap::new();
            for t in toks {
                *tf.entry(t.as_str()).or_insert(0) += 1;
            }
            for (t, f) in tf {
                postings.entry(t.to_string()).or_default().push((idx, f));
            }
        }
        Bm25Index { n, avg_len, doc_len, postings }
    }

    /// BM25 (k1=1.5, b=0.75). Deduped query terms; returns (chunk_idx, score)
    /// for chunks sharing ≥1 query term, sorted by score desc then idx asc.
    pub fn scores(&self, query_terms: &[String]) -> Vec<(usize, f32)> {
        const K1: f32 = 1.5;
        const B: f32 = 0.75;
        let avg = self.avg_len.max(1e-6);
        let mut acc: HashMap<usize, f32> = HashMap::new();
        let unique: HashSet<&String> = query_terms.iter().collect();
        for qt in unique {
            let Some(plist) = self.postings.get(qt) else { continue };
            let df = plist.len() as f32;
            let idf = (((self.n as f32 - df + 0.5) / (df + 0.5)) + 1.0).ln();
            for &(idx, f) in plist {
                let f = f as f32;
                let dl = self.doc_len[idx] as f32;
                let denom = f + K1 * (1.0 - B + B * dl / avg);
                *acc.entry(idx).or_insert(0.0) += idf * (f * (K1 + 1.0)) / denom;
            }
        }
        let mut out: Vec<(usize, f32)> = acc.into_iter().collect();
        out.sort_by(|a, b| {
            b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal).then(a.0.cmp(&b.0))
        });
        out
    }
}

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

    #[test]
    fn tokenize_lowercases_and_splits_on_nonalnum() {
        assert_eq!(tokenize("Hello, World! foo_bar"), vec!["hello", "world", "foo", "bar"]);
        assert!(tokenize("   ").is_empty());
    }

    #[test]
    fn scores_rank_matching_chunk_first() {
        let toks = vec![
            tokenize("the quick brown fox"),      // 0 — has "fox"
            tokenize("a lazy sleeping dog"),      // 1 — no query terms
            tokenize("the fox jumps over"),       // 2 — has "fox"
        ];
        let idx = Bm25Index::build(&toks);
        let s = idx.scores(&tokenize("fox"));
        let ranked: Vec<usize> = s.iter().map(|(i, _)| *i).collect();
        assert!(ranked.contains(&0) && ranked.contains(&2));
        assert!(!ranked.contains(&1)); // no shared term → not scored
    }

    #[test]
    fn rare_term_outweighs_common_term() {
        // "the" appears in all 3 chunks (low idf); "zebra" in one (high idf).
        let toks = vec![
            tokenize("the zebra runs"),   // 0 — rare "zebra"
            tokenize("the cat sleeps"),   // 1 — only common "the"
            tokenize("the dog barks"),    // 2 — only common "the"
        ];
        let idx = Bm25Index::build(&toks);
        let s = idx.scores(&tokenize("the zebra"));
        assert_eq!(s.first().unwrap().0, 0); // chunk 0 (rare match) ranks first
    }
}