hypersteeldb 0.5.4

//! Sinkhorn-OT ontology discovery — the text counterpart of structured-data discover. Extract candidate
//! terms from a corpus, embed them with the model2vec static embedder, and cluster them into a MECE
//! **facet codebook** via k-means++ prototypes + entropy-regularised optimal transport (`src/text/ot.rs`).
//! No hand-authored schema — the facets emerge from the corpus. The TUI view (`src/bin/ontology.rs`)
//! animates the Sinkhorn plan sharpening and the transport cost falling as it converges.

use crate::text::{cosine, ot};

/// Stop-words excluded from candidate terms (kept small; the embedder handles the rest).
const STOP: &[&str] = &[
    "the", "a", "an", "and", "or", "of", "to", "in", "on", "for", "is", "are", "was", "were", "be",
    "by", "as", "at", "it", "its", "this", "that", "these", "those", "with", "from", "not", "but",
    "has", "have", "had", "will", "would", "can", "may", "which", "who", "what", "when", "where",
    "how", "than", "then", "into", "over", "under", "such", "also", "more", "most", "some", "any",
    "all", "each", "per", "via", "use", "used", "using", "based", "one", "two", "our", "their",
];

/// Frequency-ranked unique candidate terms from raw text: alphanumeric words ≥3 chars, stop-words
/// dropped, capped to `max` (most frequent first).
pub fn candidate_terms(text: &str, max: usize) -> Vec<String> {
    let mut freq: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
    let mut word = String::new();
    let flush = |word: &mut String, freq: &mut std::collections::HashMap<String, usize>| {
        if word.len() >= 3 && !STOP.contains(&word.as_str()) && word.chars().any(|c| c.is_alphabetic()) {
            *freq.entry(std::mem::take(word)).or_default() += 1;
        } else {
            word.clear();
        }
    };
    for ch in text.chars() {
        if ch.is_alphanumeric() {
            word.extend(ch.to_lowercase());
        } else {
            flush(&mut word, &mut freq);
        }
    }
    flush(&mut word, &mut freq);
    let mut v: Vec<(String, usize)> = freq.into_iter().collect();
    v.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
    v.truncate(max);
    v.into_iter().map(|(t, _)| t).collect()
}

/// One discovered facet cluster.
pub struct Cluster {
    pub label: String,      // the term nearest the prototype
    pub size: usize,
    pub terms: Vec<String>, // members, most-representative first
}

/// A running ontology-discovery state: term embeddings + k-means prototypes. `assign(iters)` runs the
/// Sinkhorn plan to a given iteration count so a UI can animate convergence.
pub struct OtDiscover {
    pub terms: Vec<String>,
    embs: Vec<Vec<f32>>,
    protos: Vec<Vec<f32>>,
    pub k: usize,
    eps: f32,
}

impl OtDiscover {
    /// `terms`/`embs` are aligned (L2-normalised embeddings). Prototypes come from k-means++.
    pub fn new(terms: Vec<String>, embs: Vec<Vec<f32>>, k: usize) -> OtDiscover {
        let protos = ot::kmeans(&embs, k, 60, 1);
        OtDiscover { k: protos.len(), terms, embs, protos, eps: 0.05 }
    }

    /// Run the Sinkhorn plan to `iters` iterations → per-term cluster assignment + transport cost
    /// (falls as it converges).
    pub fn assign(&self, iters: usize) -> (Vec<usize>, f32) {
        let m: Vec<Vec<f32>> = self.embs.iter().map(|x| self.protos.iter().map(|c| 1.0 - cosine(x, c)).collect()).collect();
        let (pi, cost) = ot::sinkhorn(&m, self.eps, iters.max(1));
        let assign = pi.iter().map(|row| argmax(row)).collect();
        (assign, cost)
    }

    /// Group terms by assignment into labelled clusters (most-representative term first, as the label).
    pub fn clusters(&self, assign: &[usize], per: usize) -> Vec<Cluster> {
        let mut buckets: Vec<Vec<usize>> = vec![Vec::new(); self.k];
        for (i, &a) in assign.iter().enumerate() {
            if a < self.k {
                buckets[a].push(i);
            }
        }
        let mut out: Vec<Cluster> = buckets
            .into_iter()
            .enumerate()
            .filter(|(_, m)| !m.is_empty())
            .map(|(j, mut members)| {
                members.sort_by(|&a, &b| cosine(&self.embs[b], &self.protos[j]).partial_cmp(&cosine(&self.embs[a], &self.protos[j])).unwrap_or(std::cmp::Ordering::Equal));
                let label = self.terms[members[0]].clone();
                let terms = members.iter().take(per).map(|&i| self.terms[i].clone()).collect();
                Cluster { label, size: members.len(), terms }
            })
            .collect();
        out.sort_by(|a, b| b.size.cmp(&a.size));
        out
    }
}

fn argmax(row: &[f32]) -> usize {
    let mut a = 0;
    let mut bv = f32::NEG_INFINITY;
    for (j, &p) in row.iter().enumerate() {
        if p > bv {
            bv = p;
            a = j;
        }
    }
    a
}

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

    #[test]
    fn candidate_terms_ranks_and_filters() {
        let t = "Battery battery BATTERY range range the of a electric electric electric electric vehicle";
        let terms = candidate_terms(t, 10);
        assert_eq!(terms[0], "electric"); // strictly most frequent (4×)
        assert!(terms.contains(&"battery".to_string()));
        assert!(!terms.iter().any(|w| w == "the" || w == "of")); // stop-words dropped
    }

    #[test]
    fn clusters_partition_terms() {
        // two clearly separated directions in a tiny space → two facets
        let terms: Vec<String> = ["a", "b", "c", "d"].iter().map(|s| s.to_string()).collect();
        let embs = vec![
            vec![1.0, 0.0, 0.0],
            vec![0.98, 0.2, 0.0],
            vec![0.0, 0.0, 1.0],
            vec![0.0, 0.1, 0.99],
        ];
        let d = OtDiscover::new(terms, embs, 2);
        let (assign, _cost) = d.assign(100);
        // a,b together; c,d together
        assert_eq!(assign[0], assign[1]);
        assert_eq!(assign[2], assign[3]);
        assert_ne!(assign[0], assign[2]);
        let clusters = d.clusters(&assign, 4);
        assert_eq!(clusters.len(), 2);
        assert_eq!(clusters.iter().map(|c| c.size).sum::<usize>(), 4);
    }
}