kibble 0.1.0

//! Random-hyperplane LSH (SimHash) over dense embeddings, mirroring `minhash.rs`
//! but for cosine space: sign-bit signature -> banded candidates -> exact cosine
//! confirm -> union-find. Deterministic (fixed-seed planes).

use crate::minhash::UnionFind;
use crate::vectors::cosine;
use std::collections::{HashMap, HashSet};

const N_PLANES: usize = 64;
const BANDS: usize = 16;
const ROWS_PER_BAND: usize = 4; // BANDS * ROWS_PER_BAND == N_PLANES

fn splitmix64(state: &mut u64) -> u64 {
    *state = state.wrapping_add(0x9E3779B97F4A7C15);
    let mut z = *state;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
    z ^ (z >> 31)
}

/// `N_PLANES` deterministic hyperplanes of length `dim`, entries in [-1, 1).
fn make_planes(dim: usize) -> Vec<Vec<f32>> {
    (0..N_PLANES)
        .map(|i| {
            let mut s = 0xA5A5_A5A5_0000_0000_u64 ^ i as u64;
            (0..dim)
                .map(|_| (splitmix64(&mut s) as f64 / u64::MAX as f64 * 2.0 - 1.0) as f32)
                .collect()
        })
        .collect()
}

/// 64-bit sign signature of `vec` against `planes`.
fn signature(vec: &[f32], planes: &[Vec<f32>]) -> u64 {
    let mut sig = 0u64;
    for (i, p) in planes.iter().enumerate() {
        let n = vec.len().min(p.len());
        let mut dot = 0f32;
        for k in 0..n {
            dot += vec[k] * p[k];
        }
        if dot >= 0.0 {
            sig |= 1u64 << i;
        }
    }
    sig
}

/// Bucket `(band_index, band_bits) -> [row indices]` for the given signatures.
fn band_buckets(sigs: &[Option<u64>]) -> HashMap<(usize, u64), Vec<usize>> {
    let mut buckets: HashMap<(usize, u64), Vec<usize>> = HashMap::new();
    let mask = (1u64 << ROWS_PER_BAND) - 1;
    for (idx, s) in sigs.iter().enumerate() {
        let Some(sig) = s else { continue };
        for band in 0..BANDS {
            let key = (*sig >> (band * ROWS_PER_BAND)) & mask;
            buckets.entry((band, key)).or_default().push(idx);
        }
    }
    buckets
}

fn max_dim(sets: &[&[Vec<f32>]]) -> usize {
    sets.iter().flat_map(|s| s.iter()).map(|v| v.len()).max().unwrap_or(0)
}

/// Indices to drop from `vecs`/`texts`, keeping the longest text per near-duplicate
/// cluster (ties -> lowest index).
pub fn sim_dup_drop_indices(vecs: &[Vec<f32>], texts: &[String], threshold: f64) -> HashSet<usize> {
    let n = vecs.len();
    if n == 0 {
        return HashSet::new();
    }
    let planes = make_planes(max_dim(&[vecs]));
    let sigs: Vec<Option<u64>> = vecs
        .iter()
        .map(|v| if v.is_empty() { None } else { Some(signature(v, &planes)) })
        .collect();

    let mut uf = UnionFind::new(n);
    for members in band_buckets(&sigs).values() {
        if members.len() < 2 {
            continue;
        }
        for a in 0..members.len() {
            for b in (a + 1)..members.len() {
                let (i, j) = (members[a], members[b]);
                if uf.find(i) == uf.find(j) {
                    continue;
                }
                if cosine(&vecs[i], &vecs[j]) as f64 >= threshold {
                    uf.union(i, j);
                }
            }
        }
    }

    // Representative per component: longest text (tie -> lowest index).
    let mut rep: HashMap<usize, usize> = HashMap::new();
    for i in 0..n {
        if sigs[i].is_none() {
            continue;
        }
        let r = uf.find(i);
        match rep.get(&r) {
            None => {
                rep.insert(r, i);
            }
            Some(&cur) => {
                if texts[i].chars().count() > texts[cur].chars().count() {
                    rep.insert(r, i);
                }
            }
        }
    }
    let mut drop = HashSet::new();
    for i in 0..n {
        if sigs[i].is_none() {
            continue;
        }
        if rep[&uf.find(i)] != i {
            drop.insert(i);
        }
    }
    drop
}

/// Number of `query` vectors that have a `base` vector at cosine >= `threshold`.
pub fn sim_match_count(base: &[Vec<f32>], query: &[Vec<f32>], threshold: f64) -> usize {
    if base.is_empty() || query.is_empty() {
        return 0;
    }
    let planes = make_planes(max_dim(&[base, query]));
    let base_sigs: Vec<Option<u64>> = base
        .iter()
        .map(|v| if v.is_empty() { None } else { Some(signature(v, &planes)) })
        .collect();
    let buckets = band_buckets(&base_sigs);
    let mask = (1u64 << ROWS_PER_BAND) - 1;

    let mut count = 0;
    for q in query {
        if q.is_empty() {
            continue;
        }
        let qsig = signature(q, &planes);
        let mut seen: HashSet<usize> = HashSet::new();
        let mut matched = false;
        'bands: for band in 0..BANDS {
            let key = (qsig >> (band * ROWS_PER_BAND)) & mask;
            if let Some(cands) = buckets.get(&(band, key)) {
                for &bi in cands {
                    if !seen.insert(bi) {
                        continue;
                    }
                    if cosine(q, &base[bi]) as f64 >= threshold {
                        matched = true;
                        break 'bands;
                    }
                }
            }
        }
        if matched {
            count += 1;
        }
    }
    count
}

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

    // Two near-identical vectors + one orthogonal one.
    fn v(a: f32, b: f32) -> Vec<f32> { vec![a, b, 0.0, 0.0] }

    #[test]
    fn drops_near_identical_keeps_orthogonal() {
        let vecs = vec![v(1.0, 0.0), v(0.999, 0.001), v(0.0, 1.0)];
        let texts = vec!["short".to_string(), "a much longer answer".to_string(), "other".to_string()];
        let drop = sim_dup_drop_indices(&vecs, &texts, 0.95);
        // indices 0 and 1 collide; the LONGER text (index 1) is kept, 0 dropped.
        assert!(drop.contains(&0));
        assert!(!drop.contains(&1));
        assert!(!drop.contains(&2)); // orthogonal survives
    }

    #[test]
    fn match_count_counts_leaked_queries() {
        let base = vec![v(1.0, 0.0), v(0.0, 1.0)];
        let query = vec![v(0.999, 0.001), v(0.2, 0.98), v(-1.0, 0.0)];
        // q0 ~ base0, q1 ~ base1, q2 opposes base0 (cosine -1) -> 2 matches.
        assert_eq!(sim_match_count(&base, &query, 0.9), 2);
    }

    #[test]
    fn empty_inputs_are_safe() {
        assert!(sim_dup_drop_indices(&[], &[], 0.9).is_empty());
        assert_eq!(sim_match_count(&[], &[v(1.0, 0.0)], 0.9), 0);
    }
}