geographdb-core 0.5.4

//! Octree spatial indexing for 3D-native graph database.
//!
//! This module provides an octree implementation optimized for 3D spatial
//! queries in graph databases. Unlike traditional R-trees, octrees provide
//! optimal spatial partitioning for uniform 3D data distributions and
//! enable efficient pruning of graph traversal algorithms.

use crate::storage::data_structures::NodePoint;
use glam::Vec3;
use serde::{Deserialize, Serialize};

/// Per-query instrumentation for octree lookup complexity analysis.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct OctreeQueryStats {
    /// Number of octree nodes visited during traversal.
    pub nodes_visited: usize,
    /// Number of leaf octants scanned.
    pub leaf_nodes_scanned: usize,
    /// Number of points distance-tested against the query sphere.
    pub points_tested: usize,
    /// Number of points that matched the query predicate.
    pub points_returned: usize,
}

/// An octant in 3D space, represented by its bounding box.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct BoundingBox {
    /// Minimum corner of the bounding box.
    pub min: Vec3,
    /// Maximum corner of the bounding box.
    pub max: Vec3,
}

impl BoundingBox {
    /// Create a new bounding box from minimum and maximum corners.
    pub fn new(min: Vec3, max: Vec3) -> Self {
        Self { min, max }
    }

    /// Create a new bounding box centered at origin with given half-size.
    pub fn centered(center: Vec3, half_size: f32) -> Self {
        Self {
            min: center - Vec3::splat(half_size),
            max: center + Vec3::splat(half_size),
        }
    }

    /// Check if a point is contained within this bounding box.
    pub fn contains(&self, point: Vec3) -> bool {
        point.x >= self.min.x
            && point.x <= self.max.x
            && point.y >= self.min.y
            && point.y <= self.max.y
            && point.z >= self.min.z
            && point.z <= self.max.z
    }

    /// Check if this bounding box intersects with another bounding box.
    pub fn intersects(&self, other: &BoundingBox) -> bool {
        self.min.x <= other.max.x
            && self.max.x >= other.min.x
            && self.min.y <= other.max.y
            && self.max.y >= other.min.y
            && self.min.z <= other.max.z
            && self.max.z >= other.min.z
    }

    /// Calculate the center of this bounding box.
    pub fn center(&self) -> Vec3 {
        (self.min + self.max) * 0.5
    }

    /// Calculate the size (dimensions) of this bounding box.
    pub fn size(&self) -> Vec3 {
        self.max - self.min
    }

    /// Subdivide this bounding box into 8 octants.
    pub fn subdivide(&self) -> [BoundingBox; 8] {
        let center = self.center();
        // Calculate half-size of this bounding box
        let _half_size = (self.max - self.min) * 0.5;

        [
            // Bottom octants (z < center.z)
            BoundingBox::new(self.min, center), // Bottom-front-left
            BoundingBox::new(
                Vec3::new(center.x, self.min.y, self.min.z),
                Vec3::new(self.max.x, center.y, center.z),
            ), // Bottom-front-right
            BoundingBox::new(
                Vec3::new(self.min.x, center.y, self.min.z),
                Vec3::new(center.x, self.max.y, center.z),
            ), // Bottom-back-left
            BoundingBox::new(
                Vec3::new(center.x, center.y, self.min.z),
                Vec3::new(self.max.x, self.max.y, center.z),
            ), // Bottom-back-right
            // Top octants (z >= center.z)
            BoundingBox::new(
                Vec3::new(self.min.x, self.min.y, center.z),
                Vec3::new(center.x, center.y, self.max.z),
            ), // Top-front-left
            BoundingBox::new(
                Vec3::new(center.x, self.min.y, center.z),
                Vec3::new(self.max.x, center.y, self.max.z),
            ), // Top-front-right
            BoundingBox::new(
                Vec3::new(self.min.x, center.y, center.z),
                Vec3::new(center.x, self.max.y, self.max.z),
            ), // Top-back-left
            BoundingBox::new(center, self.max), // Top-back-right
        ]
    }

    /// Minimum squared distance from a point to any point in this bounding box.
    fn min_distance_sq(&self, point: Vec3) -> f32 {
        let closest_x = point.x.clamp(self.min.x, self.max.x);
        let closest_y = point.y.clamp(self.min.y, self.max.y);
        let closest_z = point.z.clamp(self.min.z, self.max.z);
        let closest_point = Vec3::new(closest_x, closest_y, closest_z);
        (closest_point - point).length_squared()
    }
}

/// A node in the octree structure.
#[derive(Debug, Clone)]
pub struct OctreeNode {
    /// The bounding box that defines this node's spatial extent.
    pub bounds: BoundingBox,

    /// The depth of this node in the octree (0 for root).
    pub depth: u32,

    /// The nodes contained directly in this octant (if this is a leaf).
    pub nodes: Vec<NodePoint>,

    /// Child octants (if this node has been subdivided).
    pub children: Option<[Box<OctreeNode>; 8]>,

    /// Flag indicating if this node has been subdivided.
    pub is_subdivided: bool,
}

impl OctreeNode {
    /// Create a new octree node with the given bounds and depth.
    pub fn new(bounds: BoundingBox, depth: u32) -> Self {
        Self {
            bounds,
            depth,
            nodes: Vec::new(),
            children: None,
            is_subdivided: false,
        }
    }

    /// Check if this node is a leaf (has no children).
    pub fn is_leaf(&self) -> bool {
        !self.is_subdivided
    }

    /// Get the maximum number of nodes this leaf can hold before subdivision.
    pub fn capacity(&self) -> usize {
        // Capacity increases with depth to prevent excessive subdivision
        // at deeper levels where nodes are more sparse
        (8 + self.depth as usize * 2).min(32)
    }

    /// Insert a node into this octree node.
    ///
    /// If this node is a leaf and has capacity, the node is added directly.
    /// If this node is a leaf and is at capacity, it will be subdivided and
    /// the nodes redistributed.
    /// If this node has children, the node is inserted into the appropriate child.
    pub fn insert(&mut self, node: NodePoint) -> bool {
        // Check if the node is within this octant's bounds
        let point = Vec3::new(node.x, node.y, node.z);
        let mut inserted = false;
        if self.bounds.contains(point) {
            if self.is_leaf() {
                // If we have capacity, add the node directly
                if self.nodes.len() < self.capacity() || self.depth >= 16 {
                    self.nodes.push(node);
                    inserted = true;
                } else {
                    // Otherwise, subdivide and redistribute existing nodes
                    self.subdivide();

                    // Redistribute existing nodes to children
                    let mut i = 0;
                    while i < self.nodes.len() {
                        let existing_node = self.nodes[i];
                        let redistributed = self.insert_into_children(existing_node);
                        if redistributed {
                            self.nodes.swap_remove(i);
                        } else {
                            i += 1;
                        }
                    }

                    // Insert the new node into the appropriate child
                    inserted = self.insert_into_children(node);
                }
            } else {
                // Insert the new node into the appropriate child
                inserted = self.insert_into_children(node);
            }
        }
        inserted
    }

    /// Subdivide this node into 8 children.
    fn subdivide(&mut self) {
        let child_bounds = self.bounds.subdivide();
        let child_depth = self.depth + 1;

        let children = [
            Box::new(OctreeNode::new(child_bounds[0], child_depth)),
            Box::new(OctreeNode::new(child_bounds[1], child_depth)),
            Box::new(OctreeNode::new(child_bounds[2], child_depth)),
            Box::new(OctreeNode::new(child_bounds[3], child_depth)),
            Box::new(OctreeNode::new(child_bounds[4], child_depth)),
            Box::new(OctreeNode::new(child_bounds[5], child_depth)),
            Box::new(OctreeNode::new(child_bounds[6], child_depth)),
            Box::new(OctreeNode::new(child_bounds[7], child_depth)),
        ];

        self.children = Some(children);
        self.is_subdivided = true;
    }

    /// Insert a node into the appropriate child octant.
    fn insert_into_children(&mut self, node: NodePoint) -> bool {
        if let Some(ref mut children) = self.children {
            let point = Vec3::new(node.x, node.y, node.z);
            for child in children.iter_mut() {
                if child.bounds.contains(point) {
                    return child.insert(node);
                }
            }
        }
        false
    }

    /// Find all nodes within a spherical query region.
    ///
    /// This efficiently prunes octants that don't intersect with the query sphere.
    pub fn query_sphere(&self, center: Vec3, radius: f32, results: &mut Vec<NodePoint>) {
        let mut stats = OctreeQueryStats::default();
        self.query_sphere_with_stats(center, radius, results, &mut stats);
    }

    /// Find all nodes within a spherical query region while collecting traversal stats.
    pub fn query_sphere_with_stats(
        &self,
        center: Vec3,
        radius: f32,
        results: &mut Vec<NodePoint>,
        stats: &mut OctreeQueryStats,
    ) {
        stats.nodes_visited += 1;

        // Quick rejection: if this octant doesn't intersect the query sphere, skip it
        if !self.intersects_sphere(center, radius) {
            return;
        }

        // If this is a leaf, check all nodes directly
        if self.is_leaf() {
            stats.leaf_nodes_scanned += 1;
            for node in &self.nodes {
                stats.points_tested += 1;
                let node_pos = Vec3::new(node.x, node.y, node.z);
                let distance_sq = (node_pos - center).length_squared();
                if distance_sq <= radius * radius {
                    results.push(*node);
                    stats.points_returned += 1;
                }
            }
        } else {
            // Recursively check children
            if let Some(ref children) = self.children {
                for child in children {
                    child.query_sphere_with_stats(center, radius, results, stats);
                }
            }
        }
    }

    /// Check if this octant intersects with a spherical query region.
    fn intersects_sphere(&self, center: Vec3, radius: f32) -> bool {
        // Find the closest point in the bounding box to the sphere center
        let closest_x = center.x.clamp(self.bounds.min.x, self.bounds.max.x);
        let closest_y = center.y.clamp(self.bounds.min.y, self.bounds.max.y);
        let closest_z = center.z.clamp(self.bounds.min.z, self.bounds.max.z);
        let closest_point = Vec3::new(closest_x, closest_y, closest_z);

        // Check if the closest point is within the sphere
        let distance_sq = (closest_point - center).length_squared();
        distance_sq <= radius * radius
    }

    /// Get the total number of nodes in this subtree.
    pub fn node_count(&self) -> usize {
        let mut count = self.nodes.len();
        if let Some(ref children) = self.children {
            for child in children {
                count += child.node_count();
            }
        }
        count
    }

    /// Get the total number of leaf nodes in this subtree.
    pub fn leaf_count(&self) -> usize {
        if self.is_leaf() {
            1
        } else {
            let mut count = 0;
            if let Some(ref children) = self.children {
                for child in children {
                    count += child.leaf_count();
                }
            }
            count
        }
    }

    /// Calculate the maximum depth of this subtree.
    pub fn max_depth(&self) -> u32 {
        if self.is_leaf() {
            self.depth
        } else {
            let mut max_child_depth = self.depth;
            if let Some(ref children) = self.children {
                for child in children {
                    max_child_depth = max_child_depth.max(child.max_depth());
                }
            }
            max_child_depth
        }
    }
}

/// An octree spatial index for 3D point data.
///
/// Octrees provide efficient spatial partitioning by recursively subdividing
/// 3D space into 8 octants. This makes them particularly well-suited for
/// uniform 3D data distributions and enables massive pruning during spatial queries.
#[derive(Debug, Clone)]
pub struct Octree {
    /// The root node of the octree.
    root: OctreeNode,

    /// The total number of nodes in the octree.
    node_count: usize,
}

impl Octree {
    /// Create a new octree with the given bounds.
    pub fn new(bounds: BoundingBox) -> Self {
        Self {
            root: OctreeNode::new(bounds, 0),
            node_count: 0,
        }
    }

    /// Serialize the octree to bytes.
    ///
    /// This enables persistence by converting the octree structure to a byte array
    /// that can be saved to disk and later restored.
    pub fn to_bytes(&self) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
        // Collect all nodes from the octree
        let mut nodes = Vec::with_capacity(self.node_count);
        self.collect_nodes(&self.root, &mut nodes);

        // Simple binary format: [bounds_min_x, bounds_min_y, bounds_min_z, bounds_max_x, bounds_max_y, bounds_max_z, node_count, nodes...]
        let mut bytes = Vec::new();

        // Serialize bounds
        bytes.extend_from_slice(&self.root.bounds.min.x.to_le_bytes());
        bytes.extend_from_slice(&self.root.bounds.min.y.to_le_bytes());
        bytes.extend_from_slice(&self.root.bounds.min.z.to_le_bytes());
        bytes.extend_from_slice(&self.root.bounds.max.x.to_le_bytes());
        bytes.extend_from_slice(&self.root.bounds.max.y.to_le_bytes());
        bytes.extend_from_slice(&self.root.bounds.max.z.to_le_bytes());

        // Serialize node count
        bytes.extend_from_slice(&nodes.len().to_le_bytes());

        // Serialize each node
        for node in nodes {
            bytes.extend_from_slice(&node.id.to_le_bytes());
            bytes.extend_from_slice(&node.x.to_le_bytes());
            bytes.extend_from_slice(&node.y.to_le_bytes());
            bytes.extend_from_slice(&node.z.to_le_bytes());
        }

        Ok(bytes)
    }

    /// Deserialize an octree from bytes.
    ///
    /// Restores an octree that was previously serialized with `to_bytes`.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, Box<dyn std::error::Error>> {
        use std::convert::TryInto;

        let mut offset = 0;

        // Deserialize bounds
        let min_x = f32::from_le_bytes(bytes[offset..offset + 4].try_into()?);
        offset += 4;
        let min_y = f32::from_le_bytes(bytes[offset..offset + 4].try_into()?);
        offset += 4;
        let min_z = f32::from_le_bytes(bytes[offset..offset + 4].try_into()?);
        offset += 4;
        let max_x = f32::from_le_bytes(bytes[offset..offset + 4].try_into()?);
        offset += 4;
        let max_y = f32::from_le_bytes(bytes[offset..offset + 4].try_into()?);
        offset += 4;
        let max_z = f32::from_le_bytes(bytes[offset..offset + 4].try_into()?);
        offset += 4;

        let bounds = BoundingBox::new(
            Vec3::new(min_x, min_y, min_z),
            Vec3::new(max_x, max_y, max_z),
        );

        // Deserialize node count
        let node_count = usize::from_le_bytes(bytes[offset..offset + 8].try_into()?);
        offset += 8;

        // Deserialize nodes
        let mut nodes = Vec::with_capacity(node_count);
        for _ in 0..node_count {
            let id = u64::from_le_bytes(bytes[offset..offset + 8].try_into()?);
            offset += 8;
            let x = f32::from_le_bytes(bytes[offset..offset + 4].try_into()?);
            offset += 4;
            let y = f32::from_le_bytes(bytes[offset..offset + 4].try_into()?);
            offset += 4;
            let z = f32::from_le_bytes(bytes[offset..offset + 4].try_into()?);
            offset += 4;
            nodes.push(NodePoint { id, x, y, z });
        }

        // Rebuild the octree
        let mut octree = Self::new(bounds);
        for node in nodes {
            octree.insert(node);
        }

        Ok(octree)
    }

    /// Recursively collect all nodes from the octree.
    fn collect_nodes(&self, node: &OctreeNode, output: &mut Vec<NodePoint>) {
        // Add nodes from this octant
        output.extend(&node.nodes);

        // Recursively collect from children
        if let Some(ref children) = node.children {
            for child in children.iter() {
                self.collect_nodes(child, output);
            }
        }
    }

    /// Create a new octree that encompasses all the given nodes.
    ///
    /// This automatically calculates appropriate bounds to contain all nodes
    /// with some padding.
    pub fn from_nodes(nodes: &[NodePoint]) -> Self {
        if nodes.is_empty() {
            return Self::new(BoundingBox::new(Vec3::ZERO, Vec3::ONE));
        }

        // Calculate bounds that encompass all nodes with padding
        let mut min = Vec3::new(nodes[0].x, nodes[0].y, nodes[0].z);
        let mut max = min;

        for node in nodes {
            let pos = Vec3::new(node.x, node.y, node.z);
            min = min.min(pos);
            max = max.max(pos);
        }

        // Add 10% padding
        let size = max - min;
        let padding = size * 0.1;
        let bounds = BoundingBox::new(min - padding, max + padding);

        let mut octree = Self::new(bounds);
        for &node in nodes {
            octree.insert(node);
        }
        octree
    }

    /// Insert a node into the octree.
    pub fn insert(&mut self, node: NodePoint) -> bool {
        if self.root.insert(node) {
            self.node_count += 1;
            true
        } else {
            false
        }
    }

    /// Find all nodes within a spherical query region.
    ///
    /// This efficiently prunes octants that don't intersect with the query sphere,
    /// providing massive performance improvements over brute-force scanning.
    pub fn query_sphere(&self, center: Vec3, radius: f32) -> Vec<NodePoint> {
        let (results, _stats) = self.query_sphere_with_stats(center, radius);
        results
    }

    /// Query nodes and return both results and traversal metrics.
    pub fn query_sphere_with_stats(
        &self,
        center: Vec3,
        radius: f32,
    ) -> (Vec<NodePoint>, OctreeQueryStats) {
        let mut results = Vec::new();
        let mut stats = OctreeQueryStats::default();
        self.root
            .query_sphere_with_stats(center, radius, &mut results, &mut stats);
        (results, stats)
    }

    /// Find all nodes within a distance from a given point.
    pub fn locate_within_distance(
        &self,
        center: NodePoint,
        distance_squared: f32,
    ) -> impl Iterator<Item = NodePoint> {
        let center_vec = Vec3::new(center.x, center.y, center.z);
        let radius = distance_squared.sqrt();
        self.query_sphere(center_vec, radius).into_iter()
    }

    /// Find all nodes within an axis-aligned bounding box.
    ///
    /// This is useful for rectangular region queries.
    pub fn query_aabb(&self, bounds: &BoundingBox) -> Vec<NodePoint> {
        // For simplicity, we'll use the sphere query implementation and then filter
        // In a production implementation, we'd have a dedicated AABB query method
        let center = bounds.center();
        let size = bounds.size();
        let radius = size.length() * 0.5; // Conservative estimate

        let candidates = self.query_sphere(center, radius);
        candidates
            .into_iter()
            .filter(|node| {
                let pos = Vec3::new(node.x, node.y, node.z);
                bounds.contains(pos)
            })
            .collect()
    }

    /// Find the k nearest neighbors to a query point.
    ///
    /// Returns up to k nearest nodes sorted by ascending distance.
    /// If fewer than k nodes exist in the octree, returns all nodes.
    ///
    /// Uses a best-bin-first traversal: octants are explored in order of
    /// increasing minimum distance to the query point, and an octant is
    /// skipped once the best k candidates are all closer than the octant's
    /// minimum possible distance. This is asymptotically faster than
    /// brute-force scanning for large, spatially coherent point sets.
    ///
    /// # Arguments
    /// * `center` - The query center point
    /// * `k` - Maximum number of neighbors to return
    ///
    /// # Returns
    /// Vec of (NodePoint, distance_squared) sorted by distance
    pub fn query_knn(&self, center: Vec3, k: usize) -> Vec<(NodePoint, f32)> {
        self.query_knn_with_stats(center, k).0
    }

    /// Find the k nearest neighbors with stats.
    ///
    /// Returns the k nearest nodes along with query statistics.
    pub fn query_knn_with_stats(
        &self,
        center: Vec3,
        k: usize,
    ) -> (Vec<(NodePoint, f32)>, OctreeQueryStats) {
        let mut stats = OctreeQueryStats::default();

        if k == 0 || self.is_empty() {
            return (Vec::new(), stats);
        }

        let results = query_knn_best_bin_first(&self.root, center, k, &mut stats);
        (results, stats)
    }

    /// Get the bounding box of the entire octree.
    pub fn bounds(&self) -> &BoundingBox {
        &self.root.bounds
    }

    /// Get the total number of nodes in the octree.
    pub fn len(&self) -> usize {
        self.node_count
    }

    /// Check if the octree is empty.
    pub fn is_empty(&self) -> bool {
        self.node_count == 0
    }

    /// Get statistics about the octree structure.
    pub fn statistics(&self) -> OctreeStats {
        OctreeStats {
            node_count: self.node_count,
            leaf_count: self.root.leaf_count(),
            max_depth: self.root.max_depth(),
        }
    }
}

/// Statistics about an octree's structure.
#[derive(Debug, Clone, Copy)]
pub struct OctreeStats {
    /// Total number of nodes in the octree.
    pub node_count: usize,

    /// Total number of leaf nodes in the octree.
    pub leaf_count: usize,

    /// Maximum depth of any node in the octree.
    pub max_depth: u32,
}

// -----------------------------------------------------------------------------
// Best-bin-first k-NN helper
// -----------------------------------------------------------------------------

/// Wrapper that gives `f32` a total order so it can live in a `BinaryHeap`.
/// Distances are non-negative in this module, so bitwise ordering matches
/// numeric ordering.
#[derive(Clone, Copy, PartialEq, Eq)]
struct OrdF32(u32);

impl OrdF32 {
    fn new(v: f32) -> Self {
        Self(v.to_bits())
    }
}

impl PartialOrd for OrdF32 {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for OrdF32 {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.0.cmp(&other.0)
    }
}

fn query_knn_best_bin_first(
    root: &OctreeNode,
    center: Vec3,
    k: usize,
    stats: &mut OctreeQueryStats,
) -> Vec<(NodePoint, f32)> {
    use std::cmp::Reverse;
    use std::collections::BinaryHeap;

    // Min-heap of octants to explore, ordered by minimum distance from the
    // query point to the octant bounding box.
    let mut queue: BinaryHeap<Reverse<(OrdF32, *const OctreeNode)>> = BinaryHeap::new();
    queue.push(Reverse((
        OrdF32::new(root.bounds.min_distance_sq(center)),
        root as *const OctreeNode,
    )));
    stats.nodes_visited += 1;

    // Max-heap of the best k candidates found so far, ordered by distance.
    // We keep the worst candidate on top so we can prune distant octants.
    let mut best: BinaryHeap<(OrdF32, NodePoint)> = BinaryHeap::with_capacity(k);

    while let Some(Reverse((min_dist_sq, node_ptr))) = queue.pop() {
        // Safety: pointers come only from references held by the tree, and we
        // never mutate the tree during the query.
        let node = unsafe { &*node_ptr };

        // Prune: even the closest point in this octant is farther than our
        // k-th best candidate.
        if best.len() == k && min_dist_sq > best.peek().unwrap().0 {
            break;
        }

        if node.is_leaf() {
            stats.leaf_nodes_scanned += 1;
            for point in &node.nodes {
                stats.points_tested += 1;
                let point_dist_sq =
                    (Vec3::new(point.x, point.y, point.z) - center).length_squared();
                let ord_dist = OrdF32::new(point_dist_sq);
                if best.len() < k {
                    best.push((ord_dist, *point));
                } else if ord_dist < best.peek().unwrap().0 {
                    best.pop();
                    best.push((ord_dist, *point));
                }
            }
        } else if let Some(ref children) = node.children {
            for child in children.iter() {
                queue.push(Reverse((
                    OrdF32::new(child.bounds.min_distance_sq(center)),
                    child.as_ref() as *const OctreeNode,
                )));
                stats.nodes_visited += 1;
            }
        }
    }

    stats.points_returned = best.len();
    let mut results: Vec<(NodePoint, f32)> = best
        .into_iter()
        .map(|(d, p)| (p, f32::from_bits(d.0)))
        .collect();
    results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
    results
}

#[cfg(test)]
mod tests;