dedup_mesh 0.2.0

use crate::DeDupError;
use crate::scalar::ScalarKernel;
use crate::util::Aabb;
use num_traits::AsPrimitive;
use rustc_hash::FxHashMap;
use smallvec::SmallVec;
use std::fmt::Debug;
use std::hash::Hash;

#[derive(Debug, Eq, Clone, Copy, PartialEq)]
pub enum SignedPrimInt {
    I16,
    I32,
    I64,
}

#[allow(dead_code)]
pub trait HashType<T: ScalarKernel>: num_traits::PrimInt + Debug + Hash + Eq {
    const HASH_ONE: Self;
    const HASH_NEG_ONE: Self;
    const HASH_ZERO: Self;
    fn gen_hash_key(v: [T; 3], grid_size: T) -> (Self, Self, Self);

    fn gen_hash_key_and_normalized_cell_pos(
        v: [T; 3],
        grid_size: T,
    ) -> ((Self, Self, Self), [T; 3]);
}

macro_rules! impl_hash_type {
    // Main entry point - generates all combinations
    ($(($hash_type:ty, $scalar_type:ty)),* $(,)?) => {
        $(
            impl_hash_type!(@single $hash_type, $scalar_type);
        )*
    };

    // Single implementation generator
    (@single $hash_type:ty, $scalar_type:ty) => {
        impl HashType<$scalar_type> for $hash_type {
            const HASH_ONE: Self = 1;
            const HASH_NEG_ONE: Self = -1;
            const HASH_ZERO: Self = 0;

            #[inline(always)]
            fn gen_hash_key(v: [$scalar_type; 3], grid_size: $scalar_type) -> (Self, Self, Self) {
                (
                    (v[0] / grid_size).floor() as Self,
                    (v[1] / grid_size).floor() as Self,
                    (v[2] / grid_size).floor() as Self,
                )
            }

            #[inline(always)]
            fn gen_hash_key_and_normalized_cell_pos(
                vertex: [$scalar_type; 3],
                grid_size: $scalar_type
            ) -> ((Self, Self, Self), [$scalar_type; 3]) {
                let key = Self::gen_hash_key(vertex, grid_size);

                let cell_pos = [
                    (vertex[0] - (key.0 as $scalar_type * grid_size))/grid_size,
                    (vertex[1] - (key.1 as $scalar_type * grid_size))/grid_size,
                    (vertex[2] - (key.2 as $scalar_type * grid_size))/grid_size,
                ];

                (key, cell_pos)
            }
        }
    };
}

// Generate all permutations
impl_hash_type! {
    (i16, f32),
    (i32, f32),
    (i64, f32),
    (i16, f64),
    (i32, f64),
    (i64, f64),
}

pub(crate) struct HashGrid<T: ScalarKernel, H: HashType<T>> {
    #[allow(clippy::type_complexity)]
    // Grid maps cell coordinates to list of (vertex_id, position) pairs
    // TODO: usize could now be of Index type
    grid: FxHashMap<(H, H, H), SmallVec<[(usize, [T; 3]); 1]>>,
    tolerance_sq: T,
    grid_size: T,
}

impl<T, H> HashGrid<T, H>
where
    T: ScalarKernel,
    H: HashType<T>,
    usize: AsPrimitive<T>,
{
    pub fn with_tolerance_and_capacity(tolerance: T, capacity: usize) -> Self {
        Self {
            grid: FxHashMap::with_capacity_and_hasher(capacity, Default::default()),
            tolerance_sq: tolerance * tolerance,
            grid_size: tolerance * T::TWO,
        }
    }

    /// Query a vertex and return its ID. If a vertex within tolerance exists,
    /// returns the existing vertex ID.
    pub fn query_point(&mut self, vertex: [T; 3]) -> Option<usize> {
        let (key, normalized_cell_pos) =
            H::gen_hash_key_and_normalized_cell_pos(vertex, self.grid_size);

        // Determine which neighboring cells we need to check
        let deltas: [[H; 2]; 3] = [
            if normalized_cell_pos[0] < T::HALF {
                [H::HASH_NEG_ONE, H::HASH_ZERO]
            } else {
                [H::HASH_ZERO, H::HASH_ONE]
            },
            if normalized_cell_pos[1] < T::HALF {
                [H::HASH_NEG_ONE, H::HASH_ZERO]
            } else {
                [H::HASH_ZERO, H::HASH_ONE]
            },
            if normalized_cell_pos[2] < T::HALF {
                [H::HASH_NEG_ONE, H::HASH_ZERO]
            } else {
                [H::HASH_ZERO, H::HASH_ONE]
            },
        ];

        // Check all relevant cells for existing vertices and find the closest one
        let mut closest_id = None;
        let mut closest_dist_sq = self.tolerance_sq;

        for &dx in &deltas[0] {
            for &dy in &deltas[1] {
                for &dz in &deltas[2] {
                    let check_key = (key.0 + dx, key.1 + dy, key.2 + dz);
                    if let Some(vertices) = self.grid.get(&check_key) {
                        for &(id, existing) in vertices {
                            let dist_sq = T::distance_squared(vertex, existing);
                            if dist_sq < closest_dist_sq {
                                closest_dist_sq = dist_sq;
                                closest_id = Some(id);
                            }
                        }
                    }
                }
            }
        }
        closest_id
    }

    #[inline]
    pub fn insert(&mut self, vertex: [T; 3], index: usize) {
        // Store in grid
        self.grid
            .entry(H::gen_hash_key(vertex, self.grid_size))
            .or_default()
            .push((index, vertex));
    }

    #[cfg(test)]
    pub fn conditional_insert(&mut self, vertex: [T; 3], index: usize) -> usize {
        if let Some(found_index) = self.query_point(vertex) {
            found_index
        } else {
            self.insert(vertex, index);
            index
        }
    }

    #[cfg(test)]
    pub fn debug_print(&self) {
        println!("{:?}", self.grid.len());
        println!("{:?}", self.grid)
    }
}

pub(crate) fn get_range<T>(
    aabb: Aabb<T>,
    center: [T; 3],
    longest_axis: usize,
    tolerance: T,
    allow_scaling: bool,
) -> Result<(T, SignedPrimInt), DeDupError>
where
    T: ScalarKernel,
    usize: AsPrimitive<T>,
{
    // Since we're centering the range, we only need to check the maximum absolute value
    // The range is symmetric around 0, so max_abs = max(|min|, |max|)

    let max_abs = {
        let min = aabb.min[longest_axis] - center[longest_axis];
        let max = aabb.max[longest_axis] - center[longest_axis];
        min.abs().ceil().max(max.abs().ceil())
    };
    T::optimize_grid_tolerance_and_type(max_abs, tolerance, allow_scaling)
}

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

    #[test]
    fn test_basic_deduplication() {
        let mut dedup = HashGrid::<f32, i32>::with_tolerance_and_capacity(0.1, 3);

        let id1 = dedup.conditional_insert([0.0, 0.0, 0.0], 0);
        let id2 = dedup.conditional_insert([0.05, 0.05, 0.05], 1); // Within tolerance
        let id3 = dedup.conditional_insert([1.0, 1.0, 1.0], 2); // Far away

        assert_eq!(id1, id2); // Should be deduplicated
        assert_ne!(id1, id3); // Should be different
        dedup.debug_print();
    }

    #[test]
    fn test_closest_match_selection() {
        let mut dedup = HashGrid::<f32, i32>::with_tolerance_and_capacity(0.2, 3);

        // Add two points in the same cell but different distances from query
        let id1 = dedup.conditional_insert([0.0, 0.0, 0.0], 0);
        let id2 = dedup.conditional_insert([0.3, 0.0, 0.0], 1); // Further from origin

        // Query point closer to first vertex
        let id3 = dedup.conditional_insert([0.05, 0.0, 0.0], 2);
        assert_eq!(id3, id1); // Should match closest (first) vertex

        // Query point closer to second vertex
        let id4 = dedup.conditional_insert([0.25, 0.0, 0.0], 3);
        assert_eq!(id4, id2); // Should match closest (second) vertex
    }

    #[test]
    fn test_optimized_neighbor_checking() {
        let mut dedup = HashGrid::<f32, i32>::with_tolerance_and_capacity(0.1, 3);

        // Point in center of cell - should only check current cell
        let id1 = dedup.conditional_insert([1.0, 1.0, 1.0], 0);

        // Point near edge - should check neighbor cells
        let id2 = dedup.conditional_insert([1.05, 1.05, 1.05], 1); // Within tolerance
        let id3 = dedup.conditional_insert([1.15, 1.15, 1.15], 2); // Outside tolerance

        assert_eq!(id1, id2); // Should be deduplicated
        assert_ne!(id1, id3); // Should be different
    }

    #[test]
    fn test_many_coincident_points() {
        let mut dedup = HashGrid::<f32, i32>::with_tolerance_and_capacity(0.001, 1000);

        // Add many points at the same location - this would panic kiddo
        let mut ids = Vec::new();
        for i in 0..1000 {
            ids.push(dedup.conditional_insert([1.0, 2.0, 3.0], i));
        }

        // All should have the same ID
        assert!(ids.iter().all(|&id| id == ids[0]));
    }
}