voronoid 0.0.1

use crate::bounds::BoundingBox;
use std::collections::BinaryHeap;
use std::cmp::Ordering;
use crate::algorithm::SpatialAlgorithm;

#[derive(Clone)]
struct Point {
    index: usize,
    x: f64,
    y: f64,
    z: f64,
}

/// A spatial partitioning structure based on an Octree.
///
/// This structure recursively subdivides the 3D space into eight octants.
/// It is particularly efficient for non-uniform distributions of points,
/// as it adapts the depth of the tree to the local density of points.
pub struct Algorithm3DOctree {
    bounds: BoundingBox<3>,
    capacity: usize,
    points: Vec<Point>,
    children: Option<Box<[Algorithm3DOctree; 8]>>,
}

impl Algorithm3DOctree {
    /// Creates a new `Algorithm3DOctree` with the specified bounds and capacity.
    ///
    /// # Arguments
    ///
    /// * `bounds` - The spatial boundaries of the octree.
    /// * `capacity` - The maximum number of points a leaf node can hold before subdividing.
    pub fn new(bounds: BoundingBox<3>, capacity: usize) -> Algorithm3DOctree {
        Algorithm3DOctree {
            bounds,
            capacity,
            points: Vec::new(),
            children: None,
        }
    }

    /// Inserts a point into the octree.
    ///
    /// Returns `true` if the point was successfully inserted (i.e., it is within bounds),
    /// or `false` otherwise.
    pub fn insert(&mut self, index: usize, x: f64, y: f64, z: f64) -> bool {
        if !self.contains(x, y, z) {
            return false;
        }

        if self.children.is_none() {
            if self.points.len() < self.capacity {
                self.points.push(Point { index, x, y, z });
                return true;
            }
            
            self.subdivide();
        }

        // If we have children (either existed or just created)
        if let Some(children) = &mut self.children {
            for child in children.iter_mut() {
                if child.insert(index, x, y, z) {
                    return true;
                }
            }
        }

        false
    }

    /// Clears all points from the octree, resetting it to an empty state.
    pub fn clear(&mut self) {
        self.points.clear();
        self.children = None;
    }

    fn collect_all_points(&self, points: &mut Vec<Point>) {
        points.extend_from_slice(&self.points);
        if let Some(children) = &self.children {
            for child in children.iter() {
                child.collect_all_points(points);
            }
        }
    }

    /// Returns an iterator that yields points in the octree ordered by distance from the query point.
    ///
    /// This iterator uses a priority queue to traverse the octree, ensuring that
    /// closer points (or octants containing closer points) are visited first.
    pub fn nearest_iter(&self, x: f64, y: f64, z: f64) -> NearestIterator<'_> {
        let mut queue = BinaryHeap::new();
        
        // Calculate distance to root bounds
        let d2 = dist_sq_to_bounds(x, y, z, &self.bounds);
        
        queue.push(SearchItem {
            dist_sq: d2,
            node: Some(self),
            point: None,
        });

        NearestIterator {
            queue,
            query_x: x,
            query_y: y,
            query_z: z,
        }
    }

    fn contains(&self, x: f64, y: f64, z: f64) -> bool {
        x >= self.bounds.min[0] && x <= self.bounds.max[0] &&
        y >= self.bounds.min[1] && y <= self.bounds.max[1] &&
        z >= self.bounds.min[2] && z <= self.bounds.max[2]
    }

    fn subdivide(&mut self) {
        let min_x = self.bounds.min[0];
        let min_y = self.bounds.min[1];
        let min_z = self.bounds.min[2];
        let max_x = self.bounds.max[0];
        let max_y = self.bounds.max[1];
        let max_z = self.bounds.max[2];

        let mid_x = (min_x + max_x) / 2.0;
        let mid_y = (min_y + max_y) / 2.0;
        let mid_z = (min_z + max_z) / 2.0;

        let mut children = Box::new([
            Algorithm3DOctree::new(BoundingBox::new([min_x, min_y, min_z], [mid_x, mid_y, mid_z]), self.capacity),
            Algorithm3DOctree::new(BoundingBox::new([mid_x, min_y, min_z], [max_x, mid_y, mid_z]), self.capacity),
            Algorithm3DOctree::new(BoundingBox::new([min_x, mid_y, min_z], [mid_x, max_y, mid_z]), self.capacity),
            Algorithm3DOctree::new(BoundingBox::new([mid_x, mid_y, min_z], [max_x, max_y, mid_z]), self.capacity),
            Algorithm3DOctree::new(BoundingBox::new([min_x, min_y, mid_z], [mid_x, mid_y, max_z]), self.capacity),
            Algorithm3DOctree::new(BoundingBox::new([mid_x, min_y, mid_z], [max_x, mid_y, max_z]), self.capacity),
            Algorithm3DOctree::new(BoundingBox::new([min_x, mid_y, mid_z], [mid_x, max_y, max_z]), self.capacity),
            Algorithm3DOctree::new(BoundingBox::new([mid_x, mid_y, mid_z], [max_x, max_y, max_z]), self.capacity),
        ]);

        let points = std::mem::take(&mut self.points);
        for p in points {
            for child in children.iter_mut() {
                if child.insert(p.index, p.x, p.y, p.z) {
                    break;
                }
            }
        }

        self.children = Some(children);
    }
}

impl SpatialAlgorithm<3> for Algorithm3DOctree {
    fn set_generators(&mut self, generators: &[f64], _bounds: &BoundingBox<3>) {
        self.clear();
        let count = generators.len() / 3;
        for i in 0..count {
            self.insert(i, generators[i*3], generators[i*3+1], generators[i*3+2]);
        }
    }

    fn update_generator(&mut self, index: usize, _old_pos: &[f64; 3], new_pos: &[f64; 3], _bounds: &BoundingBox<3>) {
        let mut all_points = Vec::new();
        self.collect_all_points(&mut all_points);
        
        let mut found = false;
        for p in &mut all_points {
            if p.index == index {
                p.x = new_pos[0];
                p.y = new_pos[1];
                p.z = new_pos[2];
                found = true;
                break;
            }
        }
        
        if found {
            self.set_generators_from_points(all_points);
        }
    }

    fn visit_neighbors<F>(
        &self,
        generators: &[f64],
        index: usize,
        pos: [f64; 3],
        max_dist_sq: &mut f64,
        mut visitor: F,
    ) where
        F: FnMut(usize, [f64; 3], f64) -> f64,
    {
        // Algorithm3DOctree nearest_iter yields neighbors. We can iterate until we exceed max_dist_sq?
        // The generic calculate loop checks distance, so we just need to feed candidates.
        // nearest_iter is good because it yields closest first.
        for j in self.nearest_iter(pos[0], pos[1], pos[2]) {
            if index == j { continue; }
            let ox = generators[j * 3];
            let oy = generators[j * 3 + 1];
            let oz = generators[j * 3 + 2];

            let dx = ox - pos[0];
            let dy = oy - pos[1];
            let dz = oz - pos[2];
            let dist_sq = dx * dx + dy * dy + dz * dz;

            if dist_sq > 4.0 * *max_dist_sq {
                break;
            }

            *max_dist_sq = visitor(j, [ox, oy, oz], *max_dist_sq);
        }
    }
}

impl Algorithm3DOctree {
    fn set_generators_from_points(&mut self, points: Vec<Point>) {
        self.clear();
        for p in points {
            self.insert(p.index, p.x, p.y, p.z);
        }
    }
}

struct SearchItem<'a> {
    dist_sq: f64,
    node: Option<&'a Algorithm3DOctree>,
    point: Option<&'a Point>,
}

impl<'a> PartialEq for SearchItem<'a> {
    fn eq(&self, other: &Self) -> bool {
        self.dist_sq == other.dist_sq
    }
}

impl<'a> Eq for SearchItem<'a> {}

impl<'a> PartialOrd for SearchItem<'a> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        // Reverse ordering for Min-Heap behavior
        other.dist_sq.partial_cmp(&self.dist_sq)
    }
}

impl<'a> Ord for SearchItem<'a> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.partial_cmp(other).unwrap_or(Ordering::Equal)
    }
}

/// An iterator that yields points from the octree in order of increasing distance.
pub struct NearestIterator<'a> {
    queue: BinaryHeap<SearchItem<'a>>,
    query_x: f64,
    query_y: f64,
    query_z: f64,
}

impl<'a> Iterator for NearestIterator<'a> {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(item) = self.queue.pop() {
            if let Some(point) = item.point {
                return Some(point.index);
            }

            if let Some(node) = item.node {
                if let Some(children) = &node.children {
                    for child in children.iter() {
                        let d2 = dist_sq_to_bounds(self.query_x, self.query_y, self.query_z, &child.bounds);
                        self.queue.push(SearchItem {
                            dist_sq: d2,
                            node: Some(child),
                            point: None,
                        });
                    }
                } else {
                    for p in &node.points {
                        let dx = p.x - self.query_x;
                        let dy = p.y - self.query_y;
                        let dz = p.z - self.query_z;
                        let d2 = dx * dx + dy * dy + dz * dz;
                        self.queue.push(SearchItem {
                            dist_sq: d2,
                            node: None,
                            point: Some(p),
                        });
                    }
                }
            }
        }
        None
    }
}

fn dist_sq_to_bounds(x: f64, y: f64, z: f64, b: &BoundingBox<3>) -> f64 {
    let dx = (b.min[0] - x).max(0.0).max(x - b.max[0]);
    let dy = (b.min[1] - y).max(0.0).max(y - b.max[1]);
    let dz = (b.min[2] - z).max(0.0).max(z - b.max[2]);
    dx * dx + dy * dy + dz * dz
}

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

    #[test]
    fn test_octree_insert_and_find() {
        let bounds = BoundingBox::new([0.0, 0.0, 0.0], [10.0, 10.0, 10.0]);
        let mut octree = Algorithm3DOctree::new(bounds, 1);

        octree.insert(0, 1.0, 1.0, 1.0);
        octree.insert(1, 9.0, 9.0, 9.0);
        octree.insert(2, 2.0, 2.0, 2.0);

        let mut iter = octree.nearest_iter(1.1, 1.1, 1.1);
        assert_eq!(iter.next(), Some(0));
        assert_eq!(iter.next(), Some(2));
        assert_eq!(iter.next(), Some(1));
        assert_eq!(iter.next(), None);
    }
}