seerdb 0.0.10

//! Multi-level ALEX tree implementation for improved cache locality at scale
//!
//! This module implements a hierarchical learned index structure that maintains
//! performance at 50M+ rows by using a tree of inner nodes for routing, keeping
//! the hot routing data in CPU cache.

use anyhow::{anyhow, Result};

use super::gapped_node::GappedNode;
use super::linear_model::LinearModel;

/// Multi-level ALEX tree with hierarchical structure
pub struct MultiLevelAlexTree {
    /// Root inner node
    root: Option<Box<InnerNode>>,
    /// All leaf nodes
    leaves: Vec<GappedNode>,
    /// Tree height (1 = leaves only, 2+ = has inner nodes)
    height: usize,
    /// Total number of keys
    num_keys: usize,
}

/// Inner node for routing to children
pub struct InnerNode {
    /// Learned model for predicting child position
    model: LinearModel,
    /// Child nodes (either inner nodes or leaf indices)
    children: InnerNodeChildren,
    /// Split keys between children
    split_keys: Vec<i64>,
    /// Total keys in this subtree
    _num_keys: usize,
    /// Node level (0 = parents of leaves, increases up tree)
    _level: usize,
}

/// Children of an inner node
pub enum InnerNodeChildren {
    /// Inner node children (for non-leaf parents)
    Inner(Vec<Box<InnerNode>>),
    /// Leaf node indices (for leaf parents)
    Leaves(Vec<usize>),
}

// Configuration constants
// Constants for future multi-level optimization
#[allow(dead_code)]
const MIN_FANOUT: usize = 16; // Minimum children per inner node
#[allow(dead_code)]
const MAX_FANOUT: usize = 256; // Maximum children per inner node
#[allow(dead_code)]
const BULK_BUILD_FANOUT: usize = 64; // Default fanout for bulk building

impl Default for MultiLevelAlexTree {
    fn default() -> Self {
        Self::new()
    }
}

impl MultiLevelAlexTree {
    /// Create an empty tree
    #[must_use]
    pub const fn new() -> Self {
        Self {
            root: None,
            leaves: Vec::new(),
            height: 0,
            num_keys: 0,
        }
    }

    /// Build tree from sorted data using bulk loading
    pub fn bulk_build(mut data: Vec<(i64, Vec<u8>)>) -> Result<Self> {
        if data.is_empty() {
            return Ok(Self::new());
        }

        // Ensure data is sorted
        data.sort_unstable_by_key(|(k, _)| *k);

        let num_keys = data.len();
        let height = Self::calculate_height(num_keys);

        // Build leaf nodes first
        let leaves = Self::build_leaves(&data)?;

        // Verify all keys were inserted
        let total_leaf_keys: usize = leaves
            .iter()
            .map(super::gapped_node::GappedNode::num_keys)
            .sum();
        if total_leaf_keys != num_keys {
            tracing::warn!(
                total_leaf_keys,
                num_keys,
                "Not all keys inserted into leaves during bulk build"
            );
        }

        // If only one leaf, no inner nodes needed
        if leaves.len() == 1 {
            return Ok(Self {
                root: None,
                leaves,
                height: 1,
                num_keys: total_leaf_keys,
            });
        }

        // Build inner node tree bottom-up
        let root = Self::build_inner_tree(&leaves, height - 1)?;

        Ok(Self {
            root: Some(root),
            leaves,
            height,
            num_keys: total_leaf_keys,
        })
    }

    /// Calculate appropriate tree height for given data size
    const fn calculate_height(num_keys: usize) -> usize {
        if num_keys <= 10_000 {
            1 // Single level for small data
        } else if num_keys <= 10_000_000 {
            2 // Two levels for medium data
        } else {
            3 // Three levels for large data
        }
    }

    /// Build leaf nodes from sorted data
    fn build_leaves(data: &[(i64, Vec<u8>)]) -> Result<Vec<GappedNode>> {
        let mut leaves = Vec::new();
        let keys_per_leaf = 64; // Target keys per leaf

        for chunk in data.chunks(keys_per_leaf) {
            // Create node with more capacity for gaps
            let mut node = GappedNode::new(chunk.len() * 2, 1.5);

            // Build batch for insert
            let batch: Vec<(i64, Vec<u8>)> = chunk.to_vec();

            // Use batch insert which handles sorted data better
            if !node.insert_batch(&batch)? {
                return Err(anyhow!("Failed to insert batch into leaf"));
            }

            // Retrain model after bulk insert
            node.retrain()?;

            leaves.push(node);
        }

        Ok(leaves)
    }

    /// Build inner node tree from leaves
    fn build_inner_tree(leaves: &[GappedNode], _target_height: usize) -> Result<Box<InnerNode>> {
        // Collect leaf metadata for building inner nodes
        let leaf_keys: Vec<(i64, usize)> = leaves
            .iter()
            .enumerate()
            .filter_map(|(idx, leaf)| leaf.min_key().map(|key| (key, idx)))
            .collect();

        if leaf_keys.is_empty() {
            return Err(anyhow!("No keys in leaves"));
        }

        // For now, just build a single inner node pointing to all leaves
        // This avoids the stack overflow issue
        let node = InnerNode::build_simple_root(&leaf_keys);

        Ok(Box::new(node))
    }

    /// Get value by key
    pub fn get(&self, key: i64) -> Result<Option<Vec<u8>>> {
        if self.leaves.is_empty() {
            return Ok(None);
        }

        // Find leaf containing key
        let leaf_idx = self.route_to_leaf(key)?;

        // Search within leaf
        self.leaves[leaf_idx].get(key)
    }

    /// Insert key-value pair
    pub fn insert(&mut self, key: i64, value: Vec<u8>) -> Result<()> {
        if self.leaves.is_empty() {
            // First insert - create initial leaf
            let mut leaf = GappedNode::new(64, 1.5);
            leaf.insert(key, value)?;
            self.leaves.push(leaf);
            self.num_keys = 1;
            self.height = 1;
            return Ok(());
        }

        // Route to appropriate leaf
        let leaf_idx = self.route_to_leaf(key)?;

        // Try to insert into leaf
        if self.leaves[leaf_idx].insert(key, value.clone())? {
            self.num_keys += 1;
            Ok(())
        } else {
            // Leaf is full, needs split
            self.split_leaf(leaf_idx, key, value)
        }
    }

    /// Route to the leaf that should contain the key
    fn route_to_leaf(&self, key: i64) -> Result<usize> {
        // If inner nodes exist, traverse them
        if let Some(root) = &self.root {
            return root.route_to_leaf(key);
        }

        // No inner nodes - route directly to leaves
        if self.leaves.len() == 1 {
            return Ok(0);
        }

        // Binary search on leaf boundaries
        for (i, leaf) in self.leaves.iter().enumerate() {
            if let Some(max_key) = leaf.max_key() {
                if key <= max_key {
                    return Ok(i);
                }
            }
        }

        // Key goes in last leaf
        Ok(self.leaves.len() - 1)
    }

    /// Split a leaf node when it's full
    fn split_leaf(&mut self, leaf_idx: usize, key: i64, value: Vec<u8>) -> Result<()> {
        // Split the leaf
        let (split_key, mut new_leaf) = self.leaves[leaf_idx].split()?;

        // Insert into appropriate leaf
        if key < split_key {
            self.leaves[leaf_idx].insert(key, value)?;
        } else {
            new_leaf.insert(key, value)?;
        }

        // Add new leaf
        self.leaves.push(new_leaf);
        self.num_keys += 1;

        // Update inner nodes if they exist
        if let Some(root) = &mut self.root {
            root.handle_leaf_split(leaf_idx, split_key, self.leaves.len() - 1);
        }

        Ok(())
    }

    /// Get number of keys
    #[must_use]
    pub const fn len(&self) -> usize {
        self.num_keys
    }

    /// Check if tree is empty
    #[must_use]
    pub const fn is_empty(&self) -> bool {
        self.num_keys == 0
    }

    /// Get number of leaves
    #[must_use]
    pub const fn num_leaves(&self) -> usize {
        self.leaves.len()
    }

    /// Get tree height
    #[must_use]
    pub const fn height(&self) -> usize {
        self.height
    }
}

impl InnerNode {
    /// Build a simple root node pointing to all leaves
    fn build_simple_root(leaf_keys: &[(i64, usize)]) -> Self {
        // Train linear model on leaf positions
        let mut model = LinearModel::new();
        model.train(leaf_keys);

        // Extract just the leaf indices
        let leaf_indices: Vec<usize> = leaf_keys.iter().map(|(_, idx)| *idx).collect();

        // Create split keys (first key of each leaf except the first)
        let mut split_keys = Vec::new();
        for leaf_key in leaf_keys.iter().skip(1) {
            split_keys.push(leaf_key.0);
        }

        Self {
            model,
            children: InnerNodeChildren::Leaves(leaf_indices),
            split_keys,
            _num_keys: leaf_keys.len(),
            _level: 0,
        }
    }

    /// Build inner node from leaf metadata (recursive version - currently unused)
    #[allow(dead_code)]
    fn build_from_leaves(
        leaf_keys: &[(i64, usize)],
        level: usize,
        target_height: usize,
    ) -> Result<Self> {
        // Train linear model
        let mut model = LinearModel::new();
        model.train(leaf_keys);

        // Calculate fanout
        let fanout = Self::calculate_fanout(leaf_keys.len());

        // Partition leaves into groups
        let groups = Self::partition_leaves(leaf_keys, fanout);

        // Extract split keys
        let split_keys = Self::extract_split_keys(&groups);

        // Create children
        let children = if level < target_height - 1 {
            // Need more inner node levels
            let inner_children: Result<Vec<Box<Self>>> = groups
                .into_iter()
                .map(|group| Self::build_from_leaves(&group, level + 1, target_height))
                .map(|r| r.map(Box::new))
                .collect();

            InnerNodeChildren::Inner(inner_children?)
        } else {
            // Direct leaf children
            let leaf_indices: Vec<usize> = groups
                .into_iter()
                .flat_map(|group| group.into_iter().map(|(_, idx)| idx))
                .collect();

            InnerNodeChildren::Leaves(leaf_indices)
        };

        Ok(Self {
            model,
            children,
            split_keys,
            _num_keys: leaf_keys.len(),
            _level: level,
        })
    }

    /// Calculate appropriate fanout for node
    #[allow(dead_code)]
    fn calculate_fanout(num_children: usize) -> usize {
        if num_children <= MIN_FANOUT {
            MIN_FANOUT
        } else if num_children <= BULK_BUILD_FANOUT {
            num_children
        } else {
            BULK_BUILD_FANOUT.min(MAX_FANOUT)
        }
    }

    /// Partition leaves into groups for children
    #[allow(dead_code)]
    fn partition_leaves(leaves: &[(i64, usize)], fanout: usize) -> Vec<Vec<(i64, usize)>> {
        let mut groups = Vec::new();
        let chunk_size = leaves.len().div_ceil(fanout);

        for chunk in leaves.chunks(chunk_size) {
            groups.push(chunk.to_vec());
        }

        groups
    }

    /// Extract split keys from partitioned groups
    #[allow(dead_code)]
    fn extract_split_keys(groups: &[Vec<(i64, usize)>]) -> Vec<i64> {
        let mut split_keys = Vec::new();

        for group in groups.iter().skip(1) {
            if let Some((key, _)) = group.first() {
                split_keys.push(*key);
            }
        }

        split_keys
    }

    /// Route to leaf index
    fn route_to_leaf(&self, key: i64) -> Result<usize> {
        // Use model for initial prediction
        let predicted = self.model.predict(key);

        // Find child using split keys
        let child_idx = self.find_child(key, predicted);

        match &self.children {
            InnerNodeChildren::Inner(children) => {
                // Recursively route through inner child
                children[child_idx].route_to_leaf(key)
            }
            InnerNodeChildren::Leaves(indices) => {
                // Return leaf index
                Ok(indices[child_idx])
            }
        }
    }

    /// Find child index for key
    fn find_child(&self, key: i64, _predicted: usize) -> usize {
        // Binary search on split keys for correction
        match self.split_keys.binary_search(&key) {
            Ok(idx) | Err(idx) => idx.min(self.num_children() - 1),
        }
    }

    /// Get number of children
    const fn num_children(&self) -> usize {
        match &self.children {
            InnerNodeChildren::Inner(children) => children.len(),
            InnerNodeChildren::Leaves(indices) => indices.len(),
        }
    }

    /// Handle a leaf split by updating the routing structure
    fn handle_leaf_split(&mut self, old_leaf: usize, split_key: i64, new_leaf: usize) {
        // This is simplified - full implementation would update routing
        // and potentially trigger inner node splits
        match &mut self.children {
            InnerNodeChildren::Leaves(indices) => {
                // Add new leaf to appropriate position
                let insert_pos = indices
                    .iter()
                    .position(|&idx| idx == old_leaf)
                    .unwrap_or(indices.len());

                indices.insert(insert_pos + 1, new_leaf);

                // Update split keys
                if insert_pos < self.split_keys.len() {
                    self.split_keys.insert(insert_pos, split_key);
                } else {
                    self.split_keys.push(split_key);
                }
            }
            InnerNodeChildren::Inner(_) => {
                // Would need to recursively update inner children
                // Simplified for initial implementation
            }
        }
    }
}

// Extension to GappedNode for multi-level support
impl GappedNode {
    /// Get minimum key in node
    #[must_use]
    pub fn min_key(&self) -> Option<i64> {
        let pairs = self.pairs();
        pairs.iter().map(|(k, _)| *k).min()
    }

    /// Get maximum key in node
    #[must_use]
    pub fn max_key(&self) -> Option<i64> {
        let pairs = self.pairs();
        pairs.iter().map(|(k, _)| *k).max()
    }
}

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

    #[test]
    fn test_empty_tree() {
        let tree = MultiLevelAlexTree::new();
        assert_eq!(tree.len(), 0);
        assert!(tree.is_empty());
        assert_eq!(tree.height(), 0);
    }

    #[test]
    fn test_single_insert() {
        let mut tree = MultiLevelAlexTree::new();
        tree.insert(42, vec![1, 2, 3]).unwrap();

        assert_eq!(tree.len(), 1);
        assert!(!tree.is_empty());
        assert_eq!(tree.height(), 1);

        let result = tree.get(42).unwrap();
        assert_eq!(result, Some(vec![1, 2, 3]));
    }

    #[test]
    fn test_bulk_build() {
        let data = vec![
            (1, vec![1]),
            (10, vec![10]),
            (20, vec![20]),
            (30, vec![30]),
            (40, vec![40]),
        ];

        let tree = MultiLevelAlexTree::bulk_build(data).unwrap();

        assert_eq!(tree.len(), 5);
        assert_eq!(tree.get(1).unwrap(), Some(vec![1]));
        assert_eq!(tree.get(20).unwrap(), Some(vec![20]));
        assert_eq!(tree.get(40).unwrap(), Some(vec![40]));
        assert_eq!(tree.get(100).unwrap(), None);
    }

    #[test]
    fn test_multiple_inserts() {
        let mut tree = MultiLevelAlexTree::new();

        for i in 0..100 {
            tree.insert(i, vec![i as u8]).unwrap();
        }

        assert_eq!(tree.len(), 100);

        for i in 0..100 {
            assert_eq!(tree.get(i).unwrap(), Some(vec![i as u8]));
        }
    }
}