ic-rig 0.2.0

//! Locality Sensitive Hashing (LSH) for approximate nearest-neighbour search.
//!
//! # The problem it solves
//!
//! A linear scan over N stored embeddings is O(N) — fine for hundreds of
//! documents, painful for tens of thousands. LSH reduces the search space to a
//! small *candidate set* in O(1) bucket lookups, then exact-scores only those.
//!
//! # How it works
//!
//! A random hyperplane through the origin splits the vector space in two.
//! Every vector gets a 1-bit label: which side it lands on (sign of dot product).
//! Stack `num_hyperplanes` such planes → a short binary hash per vector.
//!
//! Vectors with a *small angle* between them have a high probability of sharing
//! the same hash. That probability is: `1 - angle/π`.
//!
//! A single table has false negatives (similar vectors that happen to straddle
//! a hyperplane). Using `num_tables` independent sets of hyperplanes, each with
//! its own bucket map, a pair is returned as a candidate if it matches in *any*
//! table — driving false negatives toward zero.
//!
//! # Tuning
//!
//! | Parameter        | Effect                                          |
//! |------------------|-------------------------------------------------|
//! | `num_hyperplanes`| More bits → fewer candidates, faster scoring   |
//! | `num_tables`     | More tables → fewer false negatives, more RAM  |
//!
//! Start with `num_hyperplanes = 12` and `num_tables = 6` for most workloads.
//! Increase `num_tables` if you're missing relevant results. Increase
//! `num_hyperplanes` if the candidate set is still too large.
//!
//! # Example
//!
//! ```rust,no_run
//! use irig::vector_store::lsh::LshIndex;
//!
//! // 1536-dim vectors (text-embedding-3-small), 12 bits/hash, 6 tables
//! let mut index = LshIndex::new(1536, 12, 6, 42);
//!
//! let title_vec: Vec<f64> = vec![0.0; 1536];
//! let description_vec: Vec<f64> = vec![0.0; 1536];
//! let query_vec: Vec<f64> = vec![0.0; 1536];
//!
//! index.insert("page-1".into(), &title_vec);
//! index.insert("page-2".into(), &description_vec);
//!
//! let candidates: Vec<String> = index.query(&query_vec);
//! // exact-score only `candidates`, not all pages
//! ```

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

use crate::embeddings::{DistanceMetric, Embedding};

// ── PRNG ──────────────────────────────────────────────────────────────────────

/// Xorshift64 seeded from a caller-supplied value.
/// ICP canisters don't have `SystemTime`, so the seed comes from outside —
/// use `ic_cdk::api::time()` or a canister-global counter.
fn xorshift64(mut state: u64) -> impl FnMut() -> f32 {
    move || {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        // Map uniformly to [-1.0, 1.0]
        (state as i64 as f32) / (i64::MAX as f32)
    }
}

// ── LSH projection planes ─────────────────────────────────────────────────────

/// The random hyperplane matrix shared across all tables.
struct Hyperplanes {
    /// Flat storage: `num_tables * num_hyperplanes` unit vectors, each of
    /// length `dim`. Indexed as `[table * num_hyperplanes + plane][dim]`.
    planes: Vec<Vec<f32>>,
    num_hyperplanes: usize,
}

impl Hyperplanes {
    fn new(dim: usize, num_tables: usize, num_hyperplanes: usize, seed: u64) -> Self {
        let mut rand = xorshift64(seed | 1); // seed must be non-zero
        let total = num_tables * num_hyperplanes;
        let mut planes = Vec::with_capacity(total);

        for _ in 0..total {
            let mut plane: Vec<f32> = (0..dim).map(|_| rand()).collect();
            // Normalize so the dot product only measures direction, not magnitude.
            let norm: f32 = plane.iter().map(|x| x * x).sum::<f32>().sqrt();
            if norm > 0.0 {
                plane.iter_mut().for_each(|v| *v /= norm);
            }
            planes.push(plane);
        }

        Self { planes, num_hyperplanes }
    }

    /// Hash a vector against the hyperplanes of one table.
    /// Each hyperplane contributes 1 bit: 1 if dot ≥ 0, 0 otherwise.
    fn hash(&self, vector: &[f64], table_idx: usize) -> u64 {
        let start = table_idx * self.num_hyperplanes;
        let mut hash = 0u64;

        for (bit, plane) in self.planes[start..start + self.num_hyperplanes]
            .iter()
            .enumerate()
        {
            let dot: f32 = vector
                .iter()
                .zip(plane.iter())
                .map(|(&v, &h)| v as f32 * h)
                .sum();

            if dot >= 0.0 {
                hash |= 1 << bit;
            }
        }

        hash
    }
}

// ── LshIndex ──────────────────────────────────────────────────────────────────

/// Approximate nearest-neighbour index backed by LSH.
///
/// Insert embeddings during indexing, query during search.
/// The returned candidate IDs should then be exact-scored with cosine similarity.
pub struct LshIndex {
    planes: Hyperplanes,
    /// One `HashMap<hash → [id]>` per table.
    tables: Vec<HashMap<u64, Vec<String>>>,
    num_tables: usize,
}

impl LshIndex {
    /// Create a new index.
    ///
    /// - `dim`              — dimensionality of your embedding vectors
    /// - `num_hyperplanes`  — bits per hash (12–16 is a good starting range)
    /// - `num_tables`       — number of independent hash tables (4–8 typical)
    /// - `seed`             — PRNG seed; use `ic_cdk::api::time()` on ICP
    pub fn new(dim: usize, num_hyperplanes: usize, num_tables: usize, seed: u64) -> Self {
        Self {
            planes: Hyperplanes::new(dim, num_tables, num_hyperplanes, seed),
            tables: vec![HashMap::new(); num_tables],
            num_tables,
        }
    }

    /// Index an embedding under `id`.
    ///
    /// Call once per embedding at insert time. If a document produces multiple
    /// embeddings (title + description + keywords), insert each separately with
    /// the same `id` — the candidate set deduplicates by id anyway.
    pub fn insert(&mut self, id: String, embedding: &[f64]) {
        for table_idx in 0..self.num_tables {
            let hash = self.planes.hash(embedding, table_idx);
            self.tables[table_idx]
                .entry(hash)
                .or_default()
                .push(id.clone());
        }
    }

    /// Return candidate IDs whose hash matches the query in at least one table.
    ///
    /// This is the fast path. The caller is responsible for exact-scoring the
    /// candidates with cosine similarity and taking the top-N.
    pub fn query(&self, embedding: &[f64]) -> Vec<String> {
        let mut candidates = HashSet::new();

        for table_idx in 0..self.num_tables {
            let hash = self.planes.hash(embedding, table_idx);
            if let Some(ids) = self.tables[table_idx].get(&hash) {
                candidates.extend(ids.iter().cloned());
            }
        }

        candidates.into_iter().collect()
    }

    /// Number of distinct IDs in the index.
    pub fn len(&self) -> usize {
        // Count unique IDs across all tables (table 0 is representative).
        self.tables.first().map_or(0, |t| t.values().map(|v| v.len()).sum())
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn clear(&mut self) {
        self.tables.iter_mut().for_each(|t| t.clear());
    }

    /// Score and rank candidates for `query` against `store`, returning `(id, score)` pairs.
    ///
    /// The LSH bucket lookup narrows the field; `metric` exact-scores the survivors.
    /// Pass `None` to use the default `DistanceMetric::Cosine { normalized: false }`.
    ///
    /// Results are sorted best-first:
    /// - similarity metrics (`Cosine`, `DotProduct`) → descending
    /// - distance metrics (`Euclidean`, `Manhattan`, `Chebyshev`, `Angular`) → ascending
    ///
    /// IDs present in the candidate set but absent from `store` are silently skipped.
    pub fn search(
        &self,
        query: &Embedding,
        store: &HashMap<String, Embedding>,
        metric: Option<DistanceMetric>,
    ) -> Vec<(String, f64)> {
        let metric = metric.unwrap_or(DistanceMetric::Cosine { normalized: false });

        let mut ranked: Vec<(String, f64)> = self
            .query(&query.vec)
            .into_iter()
            .filter_map(|id| store.get(&id).map(|emb| (id, metric.score(query, emb))))
            .collect();

        if metric.higher_is_better() {
            ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        } else {
            ranked.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
        }

        ranked
    }
}