rstar 0.13.0

use crate::{Envelope, Point, RTreeObject, RTreeParams};

#[cfg(not(test))]
use alloc::vec::Vec;

#[allow(unused_imports)] // Import is required when building without std
use num_traits::Float;

/// Partitions elements into groups of clusters along a specific axis.
pub struct ClusterGroupIterator<T: RTreeObject> {
    remaining: Vec<T>,
    slab_size: usize,
    pub cluster_dimension: usize,
}

impl<T: RTreeObject> ClusterGroupIterator<T> {
    pub fn new(
        elements: Vec<T>,
        number_of_clusters_on_axis: usize,
        cluster_dimension: usize,
    ) -> Self {
        let slab_size = elements.len().div_ceil(number_of_clusters_on_axis);
        ClusterGroupIterator {
            remaining: elements,
            slab_size,
            cluster_dimension,
        }
    }
}

impl<T: RTreeObject> Iterator for ClusterGroupIterator<T> {
    type Item = Vec<T>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.remaining.len() {
            0 => None,
            len if len <= self.slab_size => {
                let mut last = ::core::mem::take(&mut self.remaining);
                // self.remaining retains its full original capacity across iterations
                // (drain doesn't shrink), so the final slab needs shrinking.
                last.shrink_to_fit();
                last.into()
            }
            len => {
                let slab_axis = self.cluster_dimension;
                let partition_point = len - self.slab_size;
                // Partition so that the slab elements end up at the tail.
                T::Envelope::partition_envelopes(slab_axis, &mut self.remaining, partition_point);
                // Drain from the end into a new Vec with exact capacity.
                // self.remaining keeps its allocation for reuse on the next iteration.
                self.remaining
                    .drain(partition_point..)
                    .collect::<Vec<_>>()
                    .into()
            }
        }
    }
}

/// Calculates the desired number of clusters on any axis
///
/// A 'cluster' refers to a set of elements that will finally form an rtree node.
pub fn calculate_number_of_clusters_on_axis<T, Params>(number_of_elements: usize) -> usize
where
    T: RTreeObject,
    Params: RTreeParams,
{
    let max_size = Params::MAX_SIZE as f32;
    // The depth of the resulting tree, assuming all leaf nodes will be filled up to MAX_SIZE
    let depth = (number_of_elements as f32).log(max_size).ceil() as i32;
    // The number of elements each subtree will hold
    let n_subtree = max_size.powi(depth - 1);
    // How many clusters will this node contain
    let number_of_clusters = (number_of_elements as f32 / n_subtree).ceil();

    let max_dimension = <T::Envelope as Envelope>::Point::DIMENSIONS as f32;
    // Try to split all clusters among all dimensions as evenly as possible by taking the nth root.
    number_of_clusters.powf(1. / max_dimension).floor() as usize
}

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

    #[test]
    fn test_cluster_group_iterator() {
        const SIZE: usize = 374;
        const NUMBER_OF_CLUSTERS_ON_AXIS: usize = 5;
        let elements: Vec<_> = (0..SIZE as i32).map(|i| [-i, -i]).collect();
        let slab_size = (elements.len()) / NUMBER_OF_CLUSTERS_ON_AXIS + 1;
        let slabs: Vec<_> =
            ClusterGroupIterator::new(elements, NUMBER_OF_CLUSTERS_ON_AXIS, 0).collect();
        assert_eq!(slabs.len(), NUMBER_OF_CLUSTERS_ON_AXIS);
        for slab in &slabs[0..slabs.len() - 1] {
            assert_eq!(slab.len(), slab_size);
        }
        let mut total_size = 0;
        let mut min_element_for_last_slab = i32::MAX;
        for slab in &slabs {
            total_size += slab.len();
            let current_min = slab.iter().min_by_key(|point| point[0]).unwrap();
            assert!(current_min[0] < min_element_for_last_slab);
            min_element_for_last_slab = current_min[0];
        }
        assert_eq!(total_size, SIZE);
    }

    /// Verify that slabs produced by ClusterGroupIterator don't retain
    /// excessive capacity from the parent Vec.
    #[test]
    fn test_cluster_group_iterator_no_excess_capacity() {
        const SIZE: usize = 10_000;
        const NUMBER_OF_CLUSTERS_ON_AXIS: usize = 5;
        let elements: Vec<_> = (0..SIZE as i32).map(|i| [-i, -i]).collect();
        let slabs: Vec<_> =
            ClusterGroupIterator::new(elements, NUMBER_OF_CLUSTERS_ON_AXIS, 0).collect();

        for (i, slab) in slabs.iter().enumerate() {
            let ratio = slab.capacity() as f64 / slab.len() as f64;
            assert!(
                ratio <= 2.0,
                "slab {i} has excessive capacity: len={}, cap={}, ratio={ratio:.1}x",
                slab.len(),
                slab.capacity(),
            );
        }
    }
}