gemath 0.1.0

//! Spatial data structures (prototype-quality).
//!
//! These are intended as broadphase helpers for game / simulation workloads.
//! They are allocation-backed (`Vec`) and therefore require `std` or `alloc`.

use crate::aabb2::Aabb2;
use crate::aabb3::Aabb3;
use crate::math;
use crate::vec2::Vec2;

#[cfg(feature = "std")]
use std::vec::Vec;
#[cfg(all(not(feature = "std"), feature = "alloc"))]
use alloc::vec::Vec;

/// A very simple 2D uniform grid that stores item indices in each cell.
///
/// - **Good for**: fairly-uniform distributions, frequent rebuilds, broadphase candidate generation.
/// - **Bad for**: extremely clustered distributions (hot cells), wildly varying object sizes.
#[derive(Clone, Debug)]
pub struct UniformGrid2<Unit: Copy = (), Space: Copy = ()> {
    origin: Vec2<Unit, Space>,
    dims: (usize, usize),
    cell_size: f32,
    cells: Vec<Vec<usize>>,
}

impl<Unit: Copy, Space: Copy> UniformGrid2<Unit, Space> {
    /// Creates a grid covering `[min, max]` (inclusive-ish), with the given cell size.
    ///
    /// `cell_size` must be > 0.
    pub fn new(min: Vec2<Unit, Space>, max: Vec2<Unit, Space>, cell_size: f32) -> Self {
        assert!(cell_size > 0.0);
        let size = max - min;
        let w = (math::ceil(size.x / cell_size).max(1.0)) as usize;
        let h = (math::ceil(size.y / cell_size).max(1.0)) as usize;
        let mut cells: Vec<Vec<usize>> = Vec::with_capacity(w * h);
        cells.resize_with(w * h, Vec::new);
        Self {
            origin: min,
            dims: (w, h),
            cell_size,
            cells,
        }
    }

    #[inline]
    pub fn dims(&self) -> (usize, usize) {
        self.dims
    }

    #[inline]
    pub fn cell_size(&self) -> f32 {
        self.cell_size
    }

    /// Clears all cell contents (but retains allocations).
    pub fn clear(&mut self) {
        for c in &mut self.cells {
            c.clear();
        }
    }

    /// Inserts an AABB index into all cells overlapped by the AABB.
    pub fn insert_aabb(&mut self, idx: usize, aabb: &Aabb2<Unit, Space>) {
        let min = aabb.min();
        let max = aabb.max();
        let (x0, y0) = self.cell_coords_clamped(min);
        let (x1, y1) = self.cell_coords_clamped(max);
        let w = self.dims.0;
        let cells = &mut self.cells;
        for y in y0..=y1 {
            for x in x0..=x1 {
                cells[y * w + x].push(idx);
            }
        }
    }

    /// Rebuilds the grid from scratch from a slice of AABBs.
    pub fn rebuild_from_aabbs(&mut self, aabbs: &[Aabb2<Unit, Space>]) {
        self.clear();
        for (i, a) in aabbs.iter().enumerate() {
            self.insert_aabb(i, a);
        }
    }

    /// Collects candidate indices for `query` into `out`.
    ///
    /// This output is **deduplicated** (sorted + dedup).
    pub fn query_aabb(&self, query: &Aabb2<Unit, Space>, out: &mut Vec<usize>) {
        out.clear();
        let min = query.min();
        let max = query.max();
        let (x0, y0) = self.cell_coords_clamped(min);
        let (x1, y1) = self.cell_coords_clamped(max);
        let w = self.dims.0;
        for y in y0..=y1 {
            for x in x0..=x1 {
                out.extend(self.cells[y * w + x].iter().copied());
            }
        }
        out.sort_unstable();
        out.dedup();
    }

    #[inline]
    fn cell_coords_clamped(&self, p: Vec2<Unit, Space>) -> (usize, usize) {
        let fx = (p.x - self.origin.x) / self.cell_size;
        let fy = (p.y - self.origin.y) / self.cell_size;

        let mut x = math::floor(fx) as isize;
        let mut y = math::floor(fy) as isize;

        if x < 0 {
            x = 0;
        }
        if y < 0 {
            y = 0;
        }
        if x >= self.dims.0 as isize {
            x = self.dims.0 as isize - 1;
        }
        if y >= self.dims.1 as isize {
            y = self.dims.1 as isize - 1;
        }
        (x as usize, y as usize)
    }
}

/// A simple AABB quadtree (prototype) for 2D broadphase.
///
/// Stores items that *do not fit fully inside any one child* at the current node.
#[derive(Clone, Debug)]
pub struct Quadtree2<Unit: Copy = (), Space: Copy = ()> {
    max_depth: u8,
    max_items_per_node: usize,
    nodes: Vec<QuadtreeNode<Unit, Space>>,
}

#[derive(Clone, Debug)]
struct QuadtreeNode<Unit: Copy, Space: Copy> {
    bounds: Aabb2<Unit, Space>,
    depth: u8,
    children: [usize; 4], // NW, NE, SW, SE (usize::MAX = none)
    items: Vec<QuadtreeItem<Unit, Space>>,
}

#[derive(Clone, Copy, Debug)]
struct QuadtreeItem<Unit: Copy, Space: Copy> {
    idx: usize,
    aabb: Aabb2<Unit, Space>,
}

const NO_CHILD: usize = usize::MAX;

impl<Unit: Copy, Space: Copy> Quadtree2<Unit, Space> {
    pub fn new(root_bounds: Aabb2<Unit, Space>, max_depth: u8, max_items_per_node: usize) -> Self {
        assert!(max_items_per_node > 0);
        let root = QuadtreeNode {
            bounds: root_bounds,
            depth: 0,
            children: [NO_CHILD; 4],
            items: Vec::new(),
        };
        let mut nodes = Vec::new();
        nodes.push(root);
        Self {
            max_depth,
            max_items_per_node,
            nodes,
        }
    }

    pub fn clear(&mut self) {
        self.nodes.truncate(1);
        self.nodes[0].children = [NO_CHILD; 4];
        self.nodes[0].items.clear();
    }

    /// Inserts an item index with its AABB.
    pub fn insert(&mut self, idx: usize, aabb: Aabb2<Unit, Space>) {
        self.insert_into_node(0, QuadtreeItem { idx, aabb });
    }

    /// Rebuilds the tree from a slice of AABBs.
    pub fn rebuild_from_aabbs(&mut self, aabbs: &[Aabb2<Unit, Space>]) {
        self.clear();
        for (i, a) in aabbs.iter().copied().enumerate() {
            self.insert(i, a);
        }
    }

    /// Queries the tree for indices whose stored AABB intersects `query`.
    pub fn query_aabb(&self, query: &Aabb2<Unit, Space>, out: &mut Vec<usize>) {
        out.clear();
        let mut stack: Vec<usize> = Vec::new();
        stack.push(0);
        while let Some(nid) = stack.pop() {
            let n = &self.nodes[nid];
            if n.bounds.intersection(query).is_none() {
                continue;
            }
            for it in &n.items {
                if it.aabb.intersection(query).is_some() {
                    out.push(it.idx);
                }
            }
            for &c in &n.children {
                if c != NO_CHILD {
                    stack.push(c);
                }
            }
        }
        out.sort_unstable();
        out.dedup();
    }

    fn insert_into_node(&mut self, node_id: usize, item: QuadtreeItem<Unit, Space>) {
        let depth = self.nodes[node_id].depth;
        if depth < self.max_depth {
            if let Some(child_slot) = self.child_slot_that_fits(node_id, &item.aabb) {
                let cid = self.ensure_child(node_id, child_slot);
                self.insert_into_node(cid, item);
                return;
            }
        }

        self.nodes[node_id].items.push(item);

        // Split if too many items and can still subdivide.
        if self.nodes[node_id].items.len() > self.max_items_per_node && depth < self.max_depth {
            self.split_node(node_id);
        }
    }

    fn split_node(&mut self, node_id: usize) {
        // Ensure children exist.
        for slot in 0..4 {
            self.ensure_child(node_id, slot);
        }

        // Repartition items.
        let mut i = 0;
        while i < self.nodes[node_id].items.len() {
            let item = self.nodes[node_id].items[i];
            if let Some(child_slot) = self.child_slot_that_fits(node_id, &item.aabb) {
                // Remove from this node and reinsert into child.
                let item = self.nodes[node_id].items.swap_remove(i);
                let cid = self.nodes[node_id].children[child_slot];
                self.insert_into_node(cid, item);
            } else {
                i += 1;
            }
        }
    }

    fn ensure_child(&mut self, node_id: usize, slot: usize) -> usize {
        let cid = self.nodes[node_id].children[slot];
        if cid != NO_CHILD {
            return cid;
        }
        let parent_bounds = self.nodes[node_id].bounds;
        let parent_depth = self.nodes[node_id].depth;
        let child_bounds = quadtree_child_bounds(parent_bounds, slot);
        let new_id = self.nodes.len();
        self.nodes.push(QuadtreeNode {
            bounds: child_bounds,
            depth: parent_depth + 1,
            children: [NO_CHILD; 4],
            items: Vec::new(),
        });
        self.nodes[node_id].children[slot] = new_id;
        new_id
    }

    fn child_slot_that_fits(&self, node_id: usize, aabb: &Aabb2<Unit, Space>) -> Option<usize> {
        let parent = &self.nodes[node_id].bounds;
        // Early out if aabb isn't even inside parent (prototype: we treat as no-fit, stays here).
        if !aabb_fully_inside(aabb, parent) {
            return None;
        }
        for slot in 0..4 {
            let cb = quadtree_child_bounds(*parent, slot);
            if aabb_fully_inside(aabb, &cb) {
                return Some(slot);
            }
        }
        None
    }
}

fn aabb_fully_inside<Unit: Copy, Space: Copy>(inner: &Aabb2<Unit, Space>, outer: &Aabb2<Unit, Space>) -> bool {
    let imin = inner.min();
    let imax = inner.max();
    let omin = outer.min();
    let omax = outer.max();
    imin.x >= omin.x && imin.y >= omin.y && imax.x <= omax.x && imax.y <= omax.y
}

fn quadtree_child_bounds<Unit: Copy, Space: Copy>(bounds: Aabb2<Unit, Space>, slot: usize) -> Aabb2<Unit, Space> {
    // slot ordering: 0=NW, 1=NE, 2=SW, 3=SE
    let min = bounds.min();
    let max = bounds.max();
    let mid = (min + max) * 0.5;
    let (cmin, cmax) = match slot {
        0 => (Vec2::new(min.x, mid.y), Vec2::new(mid.x, max.y)), // NW
        1 => (Vec2::new(mid.x, mid.y), Vec2::new(max.x, max.y)), // NE
        2 => (Vec2::new(min.x, min.y), Vec2::new(mid.x, mid.y)), // SW
        3 => (Vec2::new(mid.x, min.y), Vec2::new(max.x, mid.y)), // SE
        _ => unreachable!(),
    };
    Aabb2::from_min_max(cmin, cmax)
}

/// A minimal binary BVH for 3D AABBs (prototype).
///
/// Built by recursively splitting on the longest centroid axis at the median.
#[derive(Clone, Debug)]
pub struct Bvh3<Unit: Copy = (), Space: Copy = ()> {
    nodes: Vec<BvhNode<Unit, Space>>,
    root: usize,
}

#[derive(Clone, Copy, Debug)]
struct BvhItem<Unit: Copy, Space: Copy> {
    idx: usize,
    aabb: Aabb3<Unit, Space>,
}

#[derive(Clone, Debug)]
struct BvhNode<Unit: Copy, Space: Copy> {
    bounds: Aabb3<Unit, Space>,
    left: usize,
    right: usize,
    item: Option<usize>,
}

impl<Unit: Copy, Space: Copy> Bvh3<Unit, Space> {
    pub fn build(aabbs: &[Aabb3<Unit, Space>]) -> Self {
        assert!(!aabbs.is_empty());
        let mut items: Vec<BvhItem<Unit, Space>> = aabbs
            .iter()
            .copied()
            .enumerate()
            .map(|(i, aabb)| BvhItem { idx: i, aabb })
            .collect();

        let mut nodes: Vec<BvhNode<Unit, Space>> = Vec::new();
        let root = build_bvh_recursive(&mut nodes, &mut items[..]);
        Self { nodes, root }
    }

    pub fn query_aabb(&self, query: &Aabb3<Unit, Space>, out: &mut Vec<usize>) {
        out.clear();
        let mut stack: Vec<usize> = Vec::new();
        stack.push(self.root);
        while let Some(nid) = stack.pop() {
            let n = &self.nodes[nid];
            if n.bounds.intersection(query).is_none() {
                continue;
            }
            if let Some(idx) = n.item {
                out.push(idx);
            } else {
                stack.push(n.left);
                stack.push(n.right);
            }
        }
        out.sort_unstable();
        out.dedup();
    }
}

fn build_bvh_recursive<Unit: Copy, Space: Copy>(
    nodes: &mut Vec<BvhNode<Unit, Space>>,
    items: &mut [BvhItem<Unit, Space>],
) -> usize {
    debug_assert!(!items.is_empty());
    if items.len() == 1 {
        let idx = items[0].idx;
        let bounds = items[0].aabb;
        let nid = nodes.len();
        nodes.push(BvhNode {
            bounds,
            left: NO_CHILD,
            right: NO_CHILD,
            item: Some(idx),
        });
        return nid;
    }

    // Compute overall bounds and centroid bounds.
    let mut bounds = items[0].aabb;
    let mut cmin = items[0].aabb.center;
    let mut cmax = items[0].aabb.center;
    for it in &items[1..] {
        bounds = bounds.union(&it.aabb);
        let c = it.aabb.center;
        cmin = cmin.min(c);
        cmax = cmax.max(c);
    }
    let extent = cmax - cmin;
    let axis = if extent.x >= extent.y && extent.x >= extent.z {
        0
    } else if extent.y >= extent.z {
        1
    } else {
        2
    };

    items.sort_by(|a, b| {
        let ca = a.aabb.center;
        let cb = b.aabb.center;
        let ka = match axis {
            0 => ca.x,
            1 => ca.y,
            _ => ca.z,
        };
        let kb = match axis {
            0 => cb.x,
            1 => cb.y,
            _ => cb.z,
        };
        ka.partial_cmp(&kb).unwrap_or(core::cmp::Ordering::Equal)
    });

    let mid = items.len() / 2;
    let (left_items, right_items) = items.split_at_mut(mid);
    let left = build_bvh_recursive(nodes, left_items);
    let right = build_bvh_recursive(nodes, right_items);

    let nid = nodes.len();
    nodes.push(BvhNode {
        bounds,
        left,
        right,
        item: None,
    });
    nid
}