//! Mutable LSH tree for incremental updates.
//!
//! The tree routes items to leaves via PCA-based hyperplane splits at each
//! internal node. Leaves hold item indices and can be split when overfull
//! or merged when underfull.
//!
//! Uses `DensityClassifier::fit_raw_pca_flat()` for computing new hyperplanes
//! when splitting leaves, reusing the proven PCA-LSH machinery.

use crate::DensityClassifier;

/// A mutable LSH tree that supports insert/remove with automatic split/merge.
pub struct LshTree {
    nodes: Vec<TreeNode>,
    root: usize,
    embedding_dim: usize,
    num_bits: usize,
    min_leaf_size: usize,
    max_leaf_size: usize, // triggers split (e.g., 4 × min_leaf_size)
    seed: u64,
}

/// A node in the LSH tree. Internal nodes have hyperplanes and children;
/// leaves have item indices.
pub struct TreeNode {
    /// PCA hyperplanes for routing, flattened (num_bits × dim). None for leaves.
    pub hyperplanes: Option<Vec<f32>>,
    /// PCA eigenvalues, one per bit. None for leaves.
    pub eigenvalues: Option<Vec<f32>>,
    /// Child node indices in `LshTree.nodes[]`.
    pub children: Vec<usize>,
    /// Maps bucket_id → child index. Parallel with `children` conceptually,
    /// but stored as explicit pairs for sparse bucket spaces.
    pub bucket_id_to_child: Vec<(u64, usize)>,
    /// L2-normalized centroid of all items in this subtree.
    pub centroid: Vec<f32>,
    /// Item indices (leaf only). None for internal nodes.
    pub items: Option<Vec<u32>>,
    /// Total items in subtree.
    pub num_items: u32,
    /// Arrow batch index (leaf only, for file format linkage).
    pub batch_index: Option<u32>,
    /// Parent node index. None for root.
    pub parent: Option<usize>,
    /// Number of hyperplane bits at this node (0 for leaves).
    pub num_bits: u8,
}

impl TreeNode {
    fn new_leaf(items: Vec<u32>, centroid: Vec<f32>) -> Self {
        let num_items = items.len() as u32;
        TreeNode {
            hyperplanes: None,
            eigenvalues: None,
            children: Vec::new(),
            bucket_id_to_child: Vec::new(),
            centroid,
            items: Some(items),
            num_items,
            batch_index: None,
            parent: None,
            num_bits: 0,
        }
    }

    #[allow(dead_code)]
    fn new_internal(centroid: Vec<f32>, num_bits: u8) -> Self {
        TreeNode {
            hyperplanes: None,
            eigenvalues: None,
            children: Vec::new(),
            bucket_id_to_child: Vec::new(),
            centroid,
            items: None,
            num_items: 0,
            batch_index: None,
            parent: None,
            num_bits,
        }
    }

    /// Returns true if this node is a leaf (has items, no children).
    pub fn is_leaf(&self) -> bool {
        self.items.is_some()
    }
}

impl LshTree {
    /// Create an empty tree.
    pub fn new(
        embedding_dim: usize,
        num_bits: usize,
        min_leaf_size: usize,
        max_leaf_size: usize,
        seed: u64,
    ) -> Self {
        let root_node = TreeNode::new_leaf(Vec::new(), vec![0.0; embedding_dim]);
        LshTree {
            nodes: vec![root_node],
            root: 0,
            embedding_dim,
            num_bits,
            min_leaf_size,
            max_leaf_size,
            seed,
        }
    }

    /// Clear all nodes (used during deserialization).
    pub fn clear_nodes(&mut self) {
        self.nodes.clear();
    }

    /// Push a node (used during deserialization).
    pub fn push_node(&mut self, node: TreeNode) {
        self.nodes.push(node);
    }

    /// Set the root node index.
    pub fn set_root(&mut self, root: usize) {
        self.root = root;
    }

    /// Reconstruct parent pointers from children lists.
    pub fn rebuild_parent_pointers(&mut self) {
        // Clear all parent pointers first
        for node in &mut self.nodes {
            node.parent = None;
        }
        // Set parent for each child
        for i in 0..self.nodes.len() {
            let children: Vec<usize> = self.nodes[i].children.clone();
            for child_idx in children {
                if child_idx < self.nodes.len() {
                    self.nodes[child_idx].parent = Some(i);
                }
            }
        }
    }

    /// Build a tree from existing items. Recursively splits until all leaves
    /// are within size bounds.
    pub fn build(
        embeddings: &[f32],
        n: usize,
        embedding_dim: usize,
        num_bits: usize,
        min_leaf_size: usize,
        max_leaf_size: usize,
        seed: u64,
    ) -> Self {
        assert!(
            embeddings.len() >= n * embedding_dim,
            "embeddings too short: {} < {}",
            embeddings.len(),
            n * embedding_dim
        );

        let items: Vec<u32> = (0..n as u32).collect();
        let centroid = compute_centroid(embeddings, embedding_dim, &items);
        let root_node = TreeNode::new_leaf(items, centroid);

        let mut tree = LshTree {
            nodes: vec![root_node],
            root: 0,
            embedding_dim,
            num_bits,
            min_leaf_size,
            max_leaf_size,
            seed,
        };

        // Recursively split overful leaves
        if n > max_leaf_size {
            tree.split_leaf(0, embeddings);
        }

        tree
    }

    /// Number of nodes in the tree.
    pub fn num_nodes(&self) -> usize {
        self.nodes.len()
    }

    /// Get a reference to a node.
    pub fn node(&self, idx: usize) -> &TreeNode {
        &self.nodes[idx]
    }

    /// Root node index.
    pub fn root(&self) -> usize {
        self.root
    }

    /// Embedding dimension.
    pub fn embedding_dim(&self) -> usize {
        self.embedding_dim
    }

    /// Collect all leaf node indices.
    pub fn leaf_indices(&self) -> Vec<usize> {
        self.nodes
            .iter()
            .enumerate()
            .filter(|(_, n)| n.is_leaf())
            .map(|(i, _)| i)
            .collect()
    }

    /// Number of leaves in the tree.
    pub fn num_leaves(&self) -> usize {
        self.nodes.iter().filter(|n| n.is_leaf()).count()
    }

    /// Total items across all leaves.
    pub fn total_items(&self) -> u32 {
        self.nodes[self.root].num_items
    }

    /// Route an embedding through hyperplanes to find its destination leaf.
    /// Returns the leaf node index.
    pub fn find_leaf(&self, embedding: &[f32]) -> usize {
        assert_eq!(embedding.len(), self.embedding_dim);
        let mut node_idx = self.root;

        loop {
            let node = &self.nodes[node_idx];
            if node.is_leaf() {
                return node_idx;
            }

            let hp = node.hyperplanes.as_ref().unwrap();
            let nb = node.num_bits as usize;
            let bucket_id = hash_single(embedding, hp, nb, self.embedding_dim);

            // Find child for this bucket
            match node
                .bucket_id_to_child
                .iter()
                .find(|&&(bid, _)| bid == bucket_id)
            {
                Some(&(_, child_idx)) => {
                    node_idx = child_idx;
                }
                None => {
                    // Bucket not seen during build — route to nearest child by centroid similarity
                    node_idx = self.nearest_child(node_idx, embedding);
                }
            }
        }
    }

    /// Insert an item into the tree. Routes to the appropriate leaf and appends.
    /// If the leaf exceeds `max_leaf_size`, triggers a split.
    ///
    /// `embeddings` must contain ALL items (including the new one at index `idx`).
    pub fn insert(&mut self, idx: u32, embeddings: &[f32]) {
        let dim = self.embedding_dim;
        let emb = &embeddings[idx as usize * dim..(idx as usize + 1) * dim];
        let leaf_idx = self.find_leaf(emb);

        // Add to leaf
        self.nodes[leaf_idx].items.as_mut().unwrap().push(idx);
        self.nodes[leaf_idx].num_items += 1;

        // Update item counts up to root
        self.update_counts_up(leaf_idx);

        // Update centroid for this leaf
        let items = self.nodes[leaf_idx].items.as_ref().unwrap().clone();
        self.nodes[leaf_idx].centroid = compute_centroid(embeddings, dim, &items);

        // Split if overfull
        if items.len() > self.max_leaf_size {
            self.split_leaf(leaf_idx, embeddings);
        }
    }

    /// Remove an item from the tree. Finds the leaf containing it, removes it.
    /// If the leaf drops below `min_leaf_size`, triggers a merge with parent.
    ///
    /// Returns true if the item was found and removed.
    pub fn remove(&mut self, idx: u32, embeddings: &[f32]) -> bool {
        let dim = self.embedding_dim;
        let emb = &embeddings[idx as usize * dim..(idx as usize + 1) * dim];
        let leaf_idx = self.find_leaf(emb);

        let items = self.nodes[leaf_idx].items.as_mut().unwrap();
        let pos = items.iter().position(|&i| i == idx);
        match pos {
            None => false,
            Some(p) => {
                items.swap_remove(p);
                self.nodes[leaf_idx].num_items -= 1;

                // Update counts up
                self.update_counts_up(leaf_idx);

                // Update centroid
                let remaining = self.nodes[leaf_idx].items.as_ref().unwrap().clone();
                if remaining.is_empty() {
                    self.nodes[leaf_idx].centroid = vec![0.0; dim];
                } else {
                    self.nodes[leaf_idx].centroid =
                        compute_centroid(embeddings, dim, &remaining);
                }

                // Merge if underfull and has a parent
                if remaining.len() < self.min_leaf_size {
                    if let Some(parent_idx) = self.nodes[leaf_idx].parent {
                        self.merge_children(parent_idx, embeddings);
                    }
                }

                true
            }
        }
    }

    /// Split an overfull leaf into children using PCA-based LSH.
    fn split_leaf(&mut self, leaf_idx: usize, embeddings: &[f32]) {
        let dim = self.embedding_dim;
        let num_bits = self.num_bits;
        let items = self.nodes[leaf_idx].items.take().unwrap();
        let n = items.len();

        if n < self.min_leaf_size * 2 {
            // Too small to split meaningfully — put items back
            self.nodes[leaf_idx].items = Some(items);
            return;
        }

        // Extract embeddings for items in this leaf
        let mut leaf_data = vec![0.0f32; n * dim];
        for (local_i, &global_i) in items.iter().enumerate() {
            let src = &embeddings[global_i as usize * dim..(global_i as usize + 1) * dim];
            leaf_data[local_i * dim..(local_i + 1) * dim].copy_from_slice(src);
        }

        // Use DensityClassifier to compute PCA hyperplanes
        let mut clf = DensityClassifier::new(dim, num_bits, self.seed, 0, true);
        clf.fit_raw_pca_flat(&leaf_data, n, dim);

        let hp_matrix = match clf.get_hyperplanes() {
            Some(hp) => hp.clone(),
            None => {
                // PCA failed — put items back as leaf
                self.nodes[leaf_idx].items = Some(items);
                return;
            }
        };

        // Get bucket IDs for each item
        let bucket_ids = clf.get_bucket_ids();
        let eigenvalues: Vec<f32> = clf.get_eigenvalues().to_vec();

        // Flatten hyperplane matrix to Vec<f32>
        let rows = hp_matrix.nrows();
        let cols = hp_matrix.ncols();
        let mut hp_flat = Vec::with_capacity(rows * cols);
        for r in 0..rows {
            for c in 0..cols {
                hp_flat.push(hp_matrix[(r, c)]);
            }
        }

        // Group items by bucket
        let mut bucket_groups: std::collections::HashMap<u64, Vec<u32>> =
            std::collections::HashMap::new();
        for (local_i, &global_i) in items.iter().enumerate() {
            bucket_groups
                .entry(bucket_ids[local_i])
                .or_default()
                .push(global_i);
        }

        // Convert leaf to internal node
        self.nodes[leaf_idx].hyperplanes = Some(hp_flat);
        self.nodes[leaf_idx].eigenvalues = Some(eigenvalues);
        self.nodes[leaf_idx].num_bits = num_bits as u8;
        // items already taken via .take()

        // Create child leaves
        let mut bucket_to_child = Vec::new();
        let mut child_indices = Vec::new();

        // Sort buckets for deterministic ordering
        let mut sorted_buckets: Vec<u64> = bucket_groups.keys().cloned().collect();
        sorted_buckets.sort();

        for bucket_id in sorted_buckets {
            let child_items = bucket_groups.remove(&bucket_id).unwrap();
            let centroid = compute_centroid(embeddings, dim, &child_items);
            let child_node = TreeNode::new_leaf(child_items, centroid);
            let child_idx = self.nodes.len();
            self.nodes.push(child_node);
            self.nodes[child_idx].parent = Some(leaf_idx);
            bucket_to_child.push((bucket_id, child_idx));
            child_indices.push(child_idx);
        }

        self.nodes[leaf_idx].bucket_id_to_child = bucket_to_child;
        self.nodes[leaf_idx].children = child_indices.clone();

        // Recursively split any children that are still overfull
        for child_idx in child_indices {
            let child_size = self.nodes[child_idx].items.as_ref().unwrap().len();
            if child_size > self.max_leaf_size {
                self.split_leaf(child_idx, embeddings);
            }
        }
    }

    /// Merge all children of an internal node back into a single leaf.
    /// Used when a child becomes underfull after removal.
    fn merge_children(&mut self, parent_idx: usize, embeddings: &[f32]) {
        let dim = self.embedding_dim;

        // Collect all items from all descendant leaves
        let all_items = self.collect_items(parent_idx);

        // Convert parent back to leaf
        self.nodes[parent_idx].centroid = if all_items.is_empty() {
            vec![0.0; dim]
        } else {
            compute_centroid(embeddings, dim, &all_items)
        };

        self.nodes[parent_idx].items = Some(all_items);
        self.nodes[parent_idx].num_items =
            self.nodes[parent_idx].items.as_ref().unwrap().len() as u32;
        self.nodes[parent_idx].hyperplanes = None;
        self.nodes[parent_idx].eigenvalues = None;
        self.nodes[parent_idx].children.clear();
        self.nodes[parent_idx].bucket_id_to_child.clear();
        self.nodes[parent_idx].num_bits = 0;
        // Note: old child nodes become orphans in the nodes vec.
        // compact() would reclaim them.
    }

    /// Recursively collect all item indices under a node.
    fn collect_items(&self, node_idx: usize) -> Vec<u32> {
        let node = &self.nodes[node_idx];
        if let Some(items) = &node.items {
            return items.clone();
        }
        let mut result = Vec::new();
        for &child_idx in &node.children {
            result.extend(self.collect_items(child_idx));
        }
        result
    }

    /// Update num_items counts from a node up to the root.
    fn update_counts_up(&mut self, start: usize) {
        let mut idx = start;
        while let Some(parent) = self.nodes[idx].parent {
            let mut total = 0u32;
            for &child in &self.nodes[parent].children {
                total += self.nodes[child].num_items;
            }
            self.nodes[parent].num_items = total;
            idx = parent;
        }
    }

    /// Find the child of an internal node whose centroid is most similar to the embedding.
    fn nearest_child(&self, node_idx: usize, embedding: &[f32]) -> usize {
        let node = &self.nodes[node_idx];
        let mut best_idx = node.children[0];
        let mut best_sim = f32::NEG_INFINITY;

        for &child_idx in &node.children {
            let sim = cosine_similarity(embedding, &self.nodes[child_idx].centroid);
            if sim > best_sim {
                best_sim = sim;
                best_idx = child_idx;
            }
        }

        best_idx
    }

    /// Recompute all centroids bottom-up from leaf items.
    pub fn update_centroids(&mut self, embeddings: &[f32]) {
        self.update_centroid_recursive(self.root, embeddings);
    }

    fn update_centroid_recursive(&mut self, node_idx: usize, embeddings: &[f32]) {
        let dim = self.embedding_dim;

        if self.nodes[node_idx].is_leaf() {
            let items = self.nodes[node_idx].items.as_ref().unwrap().clone();
            self.nodes[node_idx].centroid = if items.is_empty() {
                vec![0.0; dim]
            } else {
                compute_centroid(embeddings, dim, &items)
            };
            return;
        }

        // Recurse into children first
        let children: Vec<usize> = self.nodes[node_idx].children.clone();
        for &child_idx in &children {
            self.update_centroid_recursive(child_idx, embeddings);
        }

        // Parent centroid = weighted average of child centroids
        let mut centroid = vec![0.0f32; dim];
        let mut total = 0u32;
        for &child_idx in &children {
            let w = self.nodes[child_idx].num_items;
            total += w;
            for j in 0..dim {
                centroid[j] += self.nodes[child_idx].centroid[j] * w as f32;
            }
        }

        if total > 0 {
            for val in &mut centroid {
                *val /= total as f32;
            }
            // L2-normalize
            let norm: f32 = centroid.iter().map(|x| x * x).sum::<f32>().sqrt();
            if norm > 0.0 {
                for val in &mut centroid {
                    *val /= norm;
                }
            }
        }
        self.nodes[node_idx].centroid = centroid;
    }

    /// Build a KNN graph over leaf centroids.
    /// Returns weighted adjacency list indexed by leaf position (not node index).
    /// Also returns the mapping from leaf position → node index.
    pub fn leaf_knn_graph(&self, k: usize) -> (Vec<Vec<(u32, f32)>>, Vec<usize>) {
        let leaves = self.leaf_indices();
        let dim = self.embedding_dim;
        let n = leaves.len();

        // Flatten leaf centroids
        let mut centroids = vec![0.0f32; n * dim];
        for (i, &leaf_idx) in leaves.iter().enumerate() {
            centroids[i * dim..(i + 1) * dim]
                .copy_from_slice(&self.nodes[leaf_idx].centroid);
        }

        let adj = crate::louvain::build_knn_graph(&centroids, n, dim, k);
        (adj, leaves)
    }
}

/// Hash a single embedding against hyperplanes to get a bucket ID.
fn hash_single(embedding: &[f32], hyperplanes: &[f32], num_bits: usize, dim: usize) -> u64 {
    let mut hash = 0u64;
    for bit in 0..num_bits {
        let hp_start = bit * dim;
        let dot: f32 = embedding
            .iter()
            .zip(&hyperplanes[hp_start..hp_start + dim])
            .map(|(a, b)| a * b)
            .sum();
        if dot > 0.0 {
            hash |= 1 << bit;
        }
    }
    hash
}

/// Compute L2-normalized centroid of items.
fn compute_centroid(embeddings: &[f32], dim: usize, items: &[u32]) -> Vec<f32> {
    if items.is_empty() {
        return vec![0.0; dim];
    }

    let mut centroid = vec![0.0f32; dim];
    for &idx in items {
        let start = idx as usize * dim;
        for j in 0..dim {
            centroid[j] += embeddings[start + j];
        }
    }

    let n = items.len() as f32;
    for val in &mut centroid {
        *val /= n;
    }

    let norm: f32 = centroid.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 0.0 {
        for val in &mut centroid {
            *val /= norm;
        }
    }
    centroid
}

/// Cosine similarity between two vectors (assumed L2-normalized).
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}

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

    /// Generate random unit vectors for testing.
    fn random_embeddings(n: usize, dim: usize, seed: u64) -> Vec<f32> {
        use rand::prelude::*;
        use rand_chacha::ChaCha8Rng;

        let mut rng = ChaCha8Rng::seed_from_u64(seed);
        let mut data = vec![0.0f32; n * dim];
        for i in 0..n {
            let start = i * dim;
            for j in 0..dim {
                data[start + j] = rng.gen::<f32>() - 0.5;
            }
            // L2 normalize
            let norm: f32 = data[start..start + dim]
                .iter()
                .map(|x| x * x)
                .sum::<f32>()
                .sqrt();
            if norm > 0.0 {
                for j in 0..dim {
                    data[start + j] /= norm;
                }
            }
        }
        data
    }

    #[test]
    fn test_empty_tree() {
        let tree = LshTree::new(8, 2, 4, 16, 42);
        assert_eq!(tree.num_leaves(), 1);
        assert_eq!(tree.total_items(), 0);
    }

    #[test]
    fn test_build_small() {
        let dim = 8;
        let n = 100;
        let embeddings = random_embeddings(n, dim, 42);

        let tree = LshTree::build(&embeddings, n, dim, 2, 4, 16, 42);

        // All items should be present
        assert_eq!(tree.total_items(), n as u32);

        // Should have multiple leaves (100 items, max_leaf=16)
        assert!(tree.num_leaves() > 1, "Should split into multiple leaves");

        // Each leaf should have <= max_leaf_size items
        for leaf_idx in tree.leaf_indices() {
            let items = tree.node(leaf_idx).items.as_ref().unwrap();
            assert!(
                items.len() <= 16,
                "Leaf has {} items, max is 16",
                items.len()
            );
        }

        // All items should appear exactly once across all leaves
        let mut all_items: Vec<u32> = Vec::new();
        for leaf_idx in tree.leaf_indices() {
            all_items.extend(tree.node(leaf_idx).items.as_ref().unwrap());
        }
        all_items.sort();
        let expected: Vec<u32> = (0..n as u32).collect();
        assert_eq!(all_items, expected);
    }

    #[test]
    fn test_find_leaf_routes_correctly() {
        let dim = 8;
        let n = 100;
        let embeddings = random_embeddings(n, dim, 42);

        let tree = LshTree::build(&embeddings, n, dim, 2, 4, 16, 42);

        // Every item should route back to the leaf containing it
        for item_idx in 0..n {
            let emb = &embeddings[item_idx * dim..(item_idx + 1) * dim];
            let leaf_idx = tree.find_leaf(emb);
            let items = tree.node(leaf_idx).items.as_ref().unwrap();
            assert!(
                items.contains(&(item_idx as u32)),
                "Item {item_idx} routed to leaf {leaf_idx} which doesn't contain it"
            );
        }
    }

    #[test]
    fn test_insert() {
        let dim = 8;
        let n = 50;
        let embeddings = random_embeddings(n + 10, dim, 42); // pre-allocate extras

        let mut tree = LshTree::build(&embeddings, n, dim, 2, 4, 16, 42);
        assert_eq!(tree.total_items(), 50);

        // Insert items 50-59
        for i in 50..60 {
            tree.insert(i as u32, &embeddings);
        }

        assert_eq!(tree.total_items(), 60);

        // Verify all 60 items are findable
        for i in 0..60 {
            let emb = &embeddings[i * dim..(i + 1) * dim];
            let leaf_idx = tree.find_leaf(emb);
            let items = tree.node(leaf_idx).items.as_ref().unwrap();
            assert!(items.contains(&(i as u32)));
        }
    }

    #[test]
    fn test_remove() {
        let dim = 8;
        let n = 50;
        let embeddings = random_embeddings(n, dim, 42);

        let mut tree = LshTree::build(&embeddings, n, dim, 2, 4, 16, 42);

        // Remove item 25
        assert!(tree.remove(25, &embeddings));
        assert_eq!(tree.total_items(), 49);

        // Item 25 should not be in any leaf
        let mut found = false;
        for leaf_idx in tree.leaf_indices() {
            let items = tree.node(leaf_idx).items.as_ref().unwrap();
            if items.contains(&25) {
                found = true;
            }
        }
        assert!(!found, "Item 25 should have been removed");
    }

    #[test]
    fn test_insert_triggers_split() {
        let dim = 8;
        let max_leaf = 8;
        // Start with items that fit in one leaf
        let n = max_leaf;
        let total = max_leaf + 10; // will insert 10 more
        let embeddings = random_embeddings(total, dim, 42);

        let mut tree = LshTree::build(&embeddings, n, dim, 2, 4, max_leaf, 42);
        assert_eq!(tree.num_leaves(), 1, "Should start as single leaf");

        // Insert enough to trigger split
        for i in n..total {
            tree.insert(i as u32, &embeddings);
        }

        assert!(
            tree.num_leaves() > 1,
            "Should have split after exceeding max_leaf_size"
        );
        assert_eq!(tree.total_items(), total as u32);
    }

    #[test]
    fn test_update_centroids() {
        let dim = 8;
        let n = 50;
        let embeddings = random_embeddings(n, dim, 42);

        let mut tree = LshTree::build(&embeddings, n, dim, 2, 4, 16, 42);

        // update_centroids should produce valid L2-normalized centroids
        tree.update_centroids(&embeddings);
        let root_centroid = &tree.node(tree.root()).centroid;
        let norm: f32 = root_centroid.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            (norm - 1.0).abs() < 1e-4,
            "Root centroid should be L2-normalized, got norm={norm}"
        );

        // Leaf centroids should match direct computation
        for leaf_idx in tree.leaf_indices() {
            let items = tree.node(leaf_idx).items.as_ref().unwrap();
            if items.is_empty() {
                continue;
            }
            let expected = compute_centroid(&embeddings, dim, items);
            let actual = &tree.node(leaf_idx).centroid;
            for (a, b) in expected.iter().zip(actual.iter()) {
                assert!(
                    (a - b).abs() < 1e-5,
                    "Leaf centroid mismatch: {a} vs {b}"
                );
            }
        }
    }

    #[test]
    fn test_leaf_knn_graph() {
        let dim = 8;
        let n = 100;
        let embeddings = random_embeddings(n, dim, 42);

        let tree = LshTree::build(&embeddings, n, dim, 2, 4, 16, 42);
        let (adj, leaf_map) = tree.leaf_knn_graph(5);

        assert_eq!(adj.len(), leaf_map.len());
        assert_eq!(adj.len(), tree.num_leaves());

        // Each leaf should have neighbors
        for a in &adj {
            assert!(!a.is_empty(), "Each leaf should have at least one neighbor");
        }
    }

    #[test]
    fn test_hash_single() {
        let dim = 4;
        // Simple hyperplane: [1, 0, 0, 0] — splits on first dim sign
        let hp = vec![1.0, 0.0, 0.0, 0.0];

        let emb_pos = [0.5, 0.1, 0.2, 0.3];
        let emb_neg = [-0.5, 0.1, 0.2, 0.3];

        let h_pos = hash_single(&emb_pos, &hp, 1, dim);
        let h_neg = hash_single(&emb_neg, &hp, 1, dim);

        assert_eq!(h_pos, 1); // positive dot → bit set
        assert_eq!(h_neg, 0); // negative dot → bit clear
    }
}
