leindex 2.1.1

//! TF-IDF embedder: persisted state, vocabulary building, and embedding.
use super::*;

const TFIDF_SCHEMA_VERSION: u32 = 1;

/// TF-IDF based embedding system for code content.
///
/// Produces 768-dimensional vectors by computing TF-IDF scores for the
/// top-768 tokens by IDF value, then L2-normalizing the result.
///
/// This provides meaningful cosine similarity (> 0 for related code) unlike
/// the previous hash-based approach which produced random vectors.
#[derive(Debug, Clone, Serialize, Deserialize)]
struct TfIdfPersistedState {
    #[serde(default)]
    schema_version: u32,
    vocab: Vec<String>,
    idf: Vec<f32>,
    dimension: usize,
    pdg_nodes: usize,
    pdg_edges: usize,
    #[serde(default)]
    pdg_fingerprint: String,
}

/// TF-IDF embedding provider for code search
///
/// Implements term frequency-inverse document frequency (TF-IDF) embeddings
/// for semantic code search. Uses stratified vocabulary selection across IDF
/// ranges to maximize coverage while maintaining fixed dimension (768).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TfIdfEmbedder {
    /// Ordered vocabulary (top-K tokens by IDF, K ≤ 768)
    pub(crate) vocab: Vec<String>,
    /// IDF values indexed by vocab position
    pub(crate) idf: Vec<f32>,
    /// Embedding dimension (matches existing vector index: 768)
    pub(crate) dimension: usize,
    /// PDG node count captured when persisted for staleness checks
    pub(crate) pdg_nodes: usize,
    /// PDG edge count captured when persisted for staleness checks
    pub(crate) pdg_edges: usize,
    /// PDG fingerprint captured when persisted for staleness checks
    pub(crate) pdg_fingerprint: String,
    /// word -> embedding slots, built on first use. Embedding a node then
    /// probes this once per distinct token (~50) instead of once per
    /// vocabulary entry (768). The vocabulary is fixed once the embedder is
    /// built, so the cache never goes stale.
    #[serde(skip)]
    pub(crate) slots: std::sync::OnceLock<HashMap<String, Vec<u32>>>,
}

impl TfIdfEmbedder {
    /// Build a TF-IDF embedder from a corpus of (id, content) documents.
    ///
    /// # Steps
    /// 1. Tokenize every document
    /// 2. Build document-frequency table (df\[token\] = # docs containing token)
    /// 3. Compute IDF = ln(N / df) per token, filtering extreme frequencies
    /// 4. Stratified vocabulary selection across the full IDF range (up to 768 tokens)
    #[cfg_attr(not(test), allow(dead_code))]
    pub fn build(documents: &[(String, String)]) -> Self {
        let tokenized: Vec<(String, Vec<String>)> = documents
            .iter()
            .map(|(id, content)| (id.clone(), tokenize_code(content)))
            .collect();
        Self::build_from_tokens(&tokenized)
    }

    /// Build a TF-IDF embedder from pre-tokenized documents.
    pub fn build_from_tokens(documents: &[(String, Vec<String>)]) -> Self {
        const TARGET_DIM: usize = crate::search::search::DEFAULT_EMBEDDING_DIMENSION;
        let n = documents.len();

        if n == 0 {
            return Self {
                vocab: Vec::new(),
                idf: Vec::new(),
                dimension: TARGET_DIM,
                pdg_nodes: 0,
                pdg_edges: 0,
                pdg_fingerprint: String::new(),
                slots: Default::default(),
            };
        }

        // Count document frequency per token
        let mut df: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
        let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
        for (_, tokens) in documents {
            seen.clear();
            for tok in tokens {
                if seen.insert(tok.as_str()) {
                    *df.entry(tok.to_string()).or_insert(0) += 1;
                }
            }
        }

        // Compute IDF for each token using a moderate-frequency filter.
        let n_f = n as f32;
        let min_df: usize = if n < 50 { 1 } else { (n / 1000).max(3) };
        let max_df: usize = if n < 50 { n } else { (n / 4).max(min_df + 1) };

        let mut idf_scores: Vec<(String, f32)> = df
            .into_iter()
            .filter(|(_, df_count)| *df_count >= min_df && *df_count <= max_df)
            .map(|(tok, df_count)| {
                let idf = (n_f / df_count as f32).ln();
                (tok, idf)
            })
            .collect();

        info!(
            vocab_candidates = idf_scores.len(),
            min_df,
            max_df,
            n_docs = n,
            "TF-IDF vocabulary candidates (moderate-IDF filter)"
        );

        // Stratified vocabulary selection using sort-based sampling.
        // Sort by IDF score, then sample at quantile boundaries to get
        // diverse coverage across the full IDF range.
        idf_scores.sort_by(|a, b| {
            a.1.partial_cmp(&b.1)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.0.cmp(&b.0))
        });

        let final_scores: Vec<(String, f32)> = if idf_scores.len() <= TARGET_DIM {
            // Fewer candidates than target — use all.
            idf_scores
        } else {
            // Sample at stratified quantile boundaries to get TARGET_DIM elements
            // covering the full IDF range.
            let total = idf_scores.len();
            let stride = total as f64 / TARGET_DIM as f64;
            (0..TARGET_DIM)
                .map(|i| {
                    let idx = ((i as f64 * stride) as usize).min(total - 1);
                    idf_scores[idx].clone()
                })
                .collect()
        };

        let idf_scores = final_scores;

        let vocab: Vec<String> = idf_scores.iter().map(|(t, _)| t.clone()).collect();
        let idf: Vec<f32> = idf_scores.iter().map(|(_, s)| *s).collect();

        Self {
            vocab,
            idf,
            dimension: TARGET_DIM,
            pdg_nodes: 0,
            pdg_edges: 0,
            pdg_fingerprint: String::new(),
            slots: Default::default(),
        }
    }

    /// Embed a text string to a 768-dimensional L2-normalized TF-IDF vector.
    pub fn embed(&self, text: &str) -> Vec<f32> {
        let mut vec = vec![0.0f32; self.dimension];

        if self.vocab.is_empty() {
            return vec;
        }

        // Compute term frequencies
        let tokens = tokenize_code(text);
        let total = tokens.len() as f32;
        if total == 0.0 {
            return vec;
        }

        let mut tf_map: std::collections::HashMap<&str, f32> = std::collections::HashMap::new();
        for tok in &tokens {
            *tf_map.entry(tok.as_str()).or_insert(0.0) += 1.0;
        }

        let slots = self.slots.get_or_init(|| {
            let mut slots: HashMap<String, Vec<u32>> = HashMap::with_capacity(self.vocab.len());
            for (slot, word) in self
                .vocab
                .iter()
                .enumerate()
                .take(self.dimension.min(self.idf.len()))
            {
                slots.entry(word.clone()).or_default().push(slot as u32);
            }
            slots
        });
        for (word, count) in &tf_map {
            if let Some(dims) = slots.get(*word) {
                for &slot in dims {
                    vec[slot as usize] = (count / total) * self.idf[slot as usize];
                }
            }
        }

        let magnitude: f32 = vec.iter().map(|v| v * v).sum::<f32>().sqrt();
        if magnitude > 1e-9 {
            for v in &mut vec {
                *v /= magnitude;
            }
        }

        vec
    }

    /// Embed pre-tokenized content, skipping the tokenize_code call.
    pub fn embed_tokens(&self, tokens: &[String]) -> Vec<f32> {
        let mut vec = vec![0.0f32; self.dimension];

        if self.vocab.is_empty() {
            return vec;
        }

        let total = tokens.len() as f32;
        if total == 0.0 {
            return vec;
        }

        let mut tf_map: std::collections::HashMap<&str, f32> = std::collections::HashMap::new();
        for tok in tokens {
            *tf_map.entry(tok.as_str()).or_insert(0.0) += 1.0;
        }

        let slots = self.slots.get_or_init(|| {
            let mut slots: HashMap<String, Vec<u32>> = HashMap::with_capacity(self.vocab.len());
            for (slot, word) in self
                .vocab
                .iter()
                .enumerate()
                .take(self.dimension.min(self.idf.len()))
            {
                slots.entry(word.clone()).or_default().push(slot as u32);
            }
            slots
        });
        for (word, count) in &tf_map {
            if let Some(dims) = slots.get(*word) {
                for &slot in dims {
                    vec[slot as usize] = (count / total) * self.idf[slot as usize];
                }
            }
        }

        let magnitude: f32 = vec.iter().map(|v| v * v).sum::<f32>().sqrt();
        if magnitude > 1e-9 {
            for v in &mut vec {
                *v /= magnitude;
            }
        }

        vec
    }

    fn from_persisted_state(state: TfIdfPersistedState) -> Option<Self> {
        if state.schema_version != TFIDF_SCHEMA_VERSION {
            tracing::warn!(
                "Persisted TF-IDF schema version {} != current {}; discarding",
                state.schema_version,
                TFIDF_SCHEMA_VERSION
            );
            return None;
        }
        if state.dimension != crate::search::search::DEFAULT_EMBEDDING_DIMENSION {
            tracing::warn!(
                "Persisted TF-IDF dimension {} != expected {}; discarding",
                state.dimension,
                crate::search::search::DEFAULT_EMBEDDING_DIMENSION
            );
            return None;
        }
        if state.vocab.len() != state.idf.len() {
            tracing::warn!(
                "Persisted TF-IDF vocab/idf length mismatch ({} != {}); discarding",
                state.vocab.len(),
                state.idf.len()
            );
            return None;
        }
        Some(Self {
            vocab: state.vocab,
            idf: state.idf,
            dimension: state.dimension,
            pdg_nodes: state.pdg_nodes,
            pdg_edges: state.pdg_edges,
            pdg_fingerprint: state.pdg_fingerprint,
            slots: Default::default(),
        })
    }

    fn persisted_state(&self, pdg: &ProgramDependenceGraph) -> TfIdfPersistedState {
        self.persisted_state_with_identity(
            pdg.node_count(),
            pdg.edge_count(),
            pdg_search_fingerprint(pdg),
        )
    }

    fn persisted_state_with_identity(
        &self,
        pdg_nodes: usize,
        pdg_edges: usize,
        pdg_fingerprint: String,
    ) -> TfIdfPersistedState {
        TfIdfPersistedState {
            schema_version: TFIDF_SCHEMA_VERSION,
            vocab: self.vocab.clone(),
            idf: self.idf.clone(),
            dimension: self.dimension,
            pdg_nodes,
            pdg_edges,
            pdg_fingerprint,
        }
    }

    /// Check if the embedder is fresh relative to the current PDG state
    ///
    /// Returns true if the embedder was built from the same PDG state
    /// (same node count, edge count, and fingerprint), indicating no reindex is needed.
    pub fn is_fresh(
        &self,
        pdg_node_count: usize,
        pdg_edge_count: usize,
        pdg_fingerprint: &str,
    ) -> bool {
        self.pdg_nodes == pdg_node_count
            && self.pdg_edges == pdg_edge_count
            && !pdg_fingerprint.is_empty()
            && self.pdg_fingerprint == pdg_fingerprint
    }

    /// Get the embedding dimension
    pub fn dimension(&self) -> usize {
        self.dimension
    }

    fn storage_path(project_path: &Path) -> PathBuf {
        project_path.join(".leindex").join("tfidf_embedder.bin")
    }

    /// Load a persisted TF-IDF embedder from storage
    ///
    /// Attempts to load a previously persisted embedder from the project's
    /// `.leindex/tfidf_embedder.bin` file. Returns None if the file doesn't exist.
    pub fn load_from_storage(project_path: &Path) -> Result<Option<Self>> {
        Self::load_from_artifact_path(&project_path.join(".leindex"))
    }

    /// Load a persisted TF-IDF embedder from an explicit storage directory.
    ///
    /// Generation hydration uses this path so an interrupted write to the
    /// mutable root cannot change the embedder selected by `CURRENT`.
    pub(crate) fn load_from_artifact_path(storage_path: &Path) -> Result<Option<Self>> {
        let path = storage_path.join("tfidf_embedder.bin");
        if !path.exists() {
            return Ok(None);
        }
        let bytes = std::fs::read(&path)
            .with_context(|| format!("Failed to read persisted embedder: {}", path.display()))?;
        let state: TfIdfPersistedState = bincode::deserialize(&bytes).with_context(|| {
            format!(
                "Failed to deserialize persisted embedder: {}",
                path.display()
            )
        })?;
        Ok(Self::from_persisted_state(state))
    }

    /// Persist the TF-IDF embedder to storage
    ///
    /// Serializes the embedder state (vocabulary, IDF scores, PDG identity)
    /// to the project's `.leindex/tfidf_embedder.bin` file for future loading.
    /// The freshness identity defaults to the in-memory PDG's; callers that
    /// just saved to storage pass `persisted_identity` so `is_fresh` compares
    /// against what a DB load actually reconstructs (in-memory graphs may
    /// hold duplicate node_ids the upsert collapses).
    pub fn persist_to_storage(
        &self,
        project_path: &Path,
        pdg: &ProgramDependenceGraph,
        persisted_identity: Option<(usize, usize, String)>,
    ) -> Result<()> {
        let path = Self::storage_path(project_path);
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent).with_context(|| {
                format!("Failed to create embedder directory: {}", parent.display())
            })?;
        }
        let state = match persisted_identity {
            Some((nodes, edges, fingerprint)) => {
                self.persisted_state_with_identity(nodes, edges, fingerprint)
            }
            None => self.persisted_state(pdg),
        };
        let payload = bincode::serialize(&state).context("Failed to serialize embedder")?;
        std::fs::write(&path, payload)
            .with_context(|| format!("Failed to persist embedder: {}", path.display()))
    }
}

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

    fn valid_state() -> TfIdfPersistedState {
        TfIdfPersistedState {
            schema_version: TFIDF_SCHEMA_VERSION,
            vocab: vec!["alpha".to_string()],
            idf: vec![1.0],
            dimension: crate::search::search::DEFAULT_EMBEDDING_DIMENSION,
            pdg_nodes: 1,
            pdg_edges: 0,
            pdg_fingerprint: "fp".to_string(),
        }
    }

    #[test]
    fn rejects_unknown_schema_version() {
        let mut state = valid_state();
        state.schema_version = 0;
        assert!(TfIdfEmbedder::from_persisted_state(state).is_none());
    }

    #[test]
    fn rejects_dimension_mismatch() {
        let mut state = valid_state();
        state.dimension = 512;
        assert!(TfIdfEmbedder::from_persisted_state(state).is_none());
    }

    #[test]
    fn rejects_vocab_idf_length_mismatch() {
        let mut state = valid_state();
        state.vocab.push("beta".to_string());
        assert!(TfIdfEmbedder::from_persisted_state(state).is_none());
    }

    #[test]
    fn accepts_valid_state_and_preserves_fingerprint() {
        let state = valid_state();
        let loaded = TfIdfEmbedder::from_persisted_state(state).expect("valid state loads");
        assert_eq!(
            loaded.dimension,
            crate::search::search::DEFAULT_EMBEDDING_DIMENSION
        );
        assert_eq!(loaded.pdg_fingerprint, "fp");
    }
}