dedup_mesh 0.2.0

// SPDX-License-Identifier: MIT OR Apache-2.0
// Copyright (c) 2025 lacklustr@protonmail.com https://github.com/eadf

use crate::hash_grid::{HashGrid, HashType};
use crate::scalar::ScalarKernel;
use crate::threading::{CompatibleWith, ThreadingDispatch, ThreadingKernel, compute_centroid};
use crate::util::UnsafeVob;
use crate::util::{Aabb, Array3};
use crate::{
    CheckFinite, CheckFinitePolicy, DeDupError, KeepUnused, PruneDegenerateEnum, PruneUnused,
    PruneUnusedEnum, RelaxTolerance, Scalar, SingleThreaded, ToleranceEnum, TopologyPolicy,
};
use crate::{IndexType, MultiThreaded};
use num_traits::AsPrimitive;
use rayon::iter::IntoParallelRefIterator;
use rayon::iter::ParallelIterator;
use smallvec::{SmallVec, smallvec};
use std::fmt::Debug;
use std::marker::PhantomData;
use vob::Vob;

#[derive(Debug)]
struct ShardResults<T, const AXIS: usize>
where
    T: ScalarKernel,
    usize: AsPrimitive<T>,
{
    indices: Vec<usize>,
    remap: Vec<usize>, // centroid index -> original vertex index
    centroids: Vec<[T; 3]>,
}

#[derive(Debug)]
struct Shard<T, H, const AXIS: usize>
where
    T: ScalarKernel,
    H: HashType<T>,
{
    min_axis: T,
    max_axis: T,
    pd_: PhantomData<fn(H) -> H>,
}

impl<T, H, const AXIS: usize> Shard<T, H, AXIS>
where
    T: ScalarKernel,
    H: HashType<T>,
    usize: AsPrimitive<T>,
{
    fn new(min_axis: T, max_axis: T) -> Self {
        /*println!(
            "creasing new shard. shard bounds: [{:?}, {:?}[",
            min_axis, max_axis
        );*/
        Self {
            min_axis,
            max_axis,
            pd_: PhantomData,
        }
    }

    fn deduplicate_shard(
        &self,
        tolerance: T,
        input_vertices: &[impl Into<[T; 3]> + Clone + Sync],
        aabb_center: [T; 3],
        used_vertices: &Option<Vob>,
        num_shards: usize,
    ) -> Option<ShardResults<T, AXIS>> {
        //let tolerance_squared = tolerance * tolerance;
        let min_overlap = self.min_axis - tolerance * 2.as_();
        let max_overlap = self.max_axis + tolerance * 2.as_();

        let mut spatial_index = HashGrid::<T, H>::with_tolerance_and_capacity(
            tolerance,
            input_vertices.len() / num_shards,
        );
        let mut clusters: Vec<SmallVec<[usize; 1]>> =
            Vec::with_capacity((input_vertices.len() / num_shards) * 2);

        let mut process_vertex = |vertex_index: usize, point: [T; 3]| {
            if point[AXIS] >= max_overlap || point[AXIS] < min_overlap {
                return;
            }

            if let Some(nearest) = spatial_index.query_point(point) {
                /*println!(
                    "processing vertex {point:?}:{vertex_index} add to existing cluster, shard bounds: [{:?}, {:?}[",
                    self.min_axis, self.max_axis
                );*/
                // Add to existing cluster
                clusters[nearest].push(vertex_index);
            } else {
                // Create new cluster
                /*println!(
                    "processing vertex {point:?}:{vertex_index} create new cluster, shard bounds: [{:?}, {:?}[",
                    self.min_axis, self.max_axis
                );*/
                let cluster_id = clusters.len();
                clusters.push(smallvec![vertex_index]);
                spatial_index.insert(point, cluster_id);
            }
        };

        if let Some(used_vertices) = used_vertices {
            for (vertex_index, point) in input_vertices.iter().enumerate().filter_map(|(i, v)| {
                if used_vertices.ᚦget(i) {
                    let v = v.clone().into().sub(aabb_center);
                    if v[AXIS] < max_overlap && v[AXIS] >= min_overlap {
                        Some((i, v))
                    } else {
                        None
                    }
                } else {
                    None
                }
            }) {
                process_vertex(vertex_index, point);
            }
        } else {
            for (vertex_index, point) in input_vertices.iter().enumerate().filter_map(|(i, v)| {
                let v = v.clone().into().sub(aabb_center);
                if v[AXIS] < max_overlap && v[AXIS] >= min_overlap {
                    Some((i, v))
                } else {
                    None
                }
            }) {
                process_vertex(vertex_index, point);
            }
        };

        let aabb_center_axis = aabb_center[AXIS];
        if clusters.is_empty() {
            //println!(
            //    "Shard shardmin:{:?} shardmax:{:?} was empty",
            //    self.min_axis, self.max_axis
            //);
            None
        } else {
            let mut centroid_id = 0_usize;
            // Compute representatives for all clusters
            let shard = clusters
                .into_iter()
                .filter_map(|participants| {
                    if let Some(centroid) =
                        compute_centroid::<T, AXIS>(&participants, input_vertices)
                    {
                        let centroid_axis = centroid[AXIS] - aabb_center_axis;
                        //println!("checking centroid original: {:?}, centroid_axis (offset): {:?}, participants:{participants:?} shard bounds: [{:?}, {:?}[", centroid, centroid_axis, self.min_axis, self.max_axis);
                        if centroid_axis >= self.min_axis && centroid_axis < self.max_axis {
                            //println!("centroid was kept. centroid_id:{centroid_id:?} participants:{participants:?}");
                            let rv = Some((
                                participants.to_vec(),
                                vec![centroid_id; participants.len()],
                                centroid,
                            ));
                            centroid_id += 1;
                            return rv;
                        }
                    }
                    None
                })
                .fold(
                    (Vec::new(), Vec::new(), Vec::new()),
                    |(mut acc_parts, mut acc_ids, mut acc_centroids): (
                        Vec<usize>,
                        Vec<usize>,
                        Vec<[T; 3]>,
                    ),
                     (parts, ids, centroid)| {
                        acc_parts.extend(parts);
                        acc_ids.extend(ids);
                        acc_centroids.push(centroid);
                        (acc_parts, acc_ids, acc_centroids)
                    },
                );
            /*if !shard.0.is_empty() {
                println!(
                    "Shard indices:{:?} remap{:?} centroids:{:?} shard bounds: [{:?}, {:?}[",
                    shard.0, shard.1, shard.2, self.min_axis, self.max_axis
                );
            } else {
                //println!(
                //    "Shard shardmin:{:?} shardmax:{:?} was empty",
                //    self.min_axis, self.max_axis
                //);
            }*/
            (!shard.0.is_empty()).then_some(ShardResults {
                indices: shard.0,
                remap: shard.1,
                centroids: shard.2,
            })
        }
    }
}

impl CompatibleWith<CheckFinite> for MultiThreaded {}

impl ThreadingDispatch for MultiThreaded {
    /// This dispatch calls SingleThreaded if there is too little data to build the shards.
    /// If there is enough data it calls itself.
    fn dedup_dispatch<T, Index, Vout, Topology>(
        vertices: &[impl Into<[T; 3]> + Clone + Sync],
        indices: &[Index],
        tolerance: T,
        prune_unused: PruneUnusedEnum,
        prune_degenerate: PruneDegenerateEnum,
        tolerance_policy: ToleranceEnum,
    ) -> Result<(Vec<Vout>, Vec<Index>), DeDupError>
    where
        T: Scalar,
        Index: IndexType,
        Topology: TopologyPolicy,
        usize: AsPrimitive<T>,
        Vout: Into<[T; 3]> + From<[T; 3]> + Clone + Sync,
    {
        // TODO: point cloud should allow  indices.len() < Topology::INDICES_MODULUS for MT
        if vertices.len() < Topology::INDICES_MODULUS || indices.len() < Topology::INDICES_MODULUS {
            // let the single threaded mode deal with this
            let used_vertices = match prune_unused {
                PruneUnused => SingleThreaded::get_unused_vertices::<T, Index>(vertices, indices)?,
                KeepUnused => None,
            };

            T::dedup_with_optimal_hash::<Index, Vout, SingleThreaded, Topology>(
                vertices,
                indices,
                tolerance,
                prune_degenerate,
                used_vertices,
                tolerance_policy == RelaxTolerance,
            )
        } else {
            let used_vertices = match prune_unused {
                PruneUnused => Self::get_unused_vertices::<T, Index>(vertices, indices)?,
                KeepUnused => None,
            };

            T::dedup_with_optimal_hash::<Index, Vout, Self, Topology>(
                vertices,
                indices,
                tolerance,
                prune_degenerate,
                used_vertices,
                tolerance_policy == RelaxTolerance,
            )
        }
    }

    fn dedup_exact_dispatch<T, Index, Vout, Topology, CheckFinite>(
        _vertices: &[impl Into<[T; 3]> + Clone + Sync],
        _indices: &[Index],
        _prune_unused: PruneUnusedEnum,
        _prune_degenerate: PruneDegenerateEnum,
    ) -> Result<(Vec<Vout>, Vec<Index>), DeDupError>
    where
        T: Scalar,
        Index: IndexType,
        Vout: Into<[T; 3]> + From<[T; 3]> + Clone + Sync,
        Topology: TopologyPolicy,
        CheckFinite: CheckFinitePolicy,
    {
        Err(DeDupError("Not implemented".to_string()))
    }
}

#[cfg(feature = "parallel")]
impl ThreadingKernel for MultiThreaded {
    fn dedup_vertices<T, Index, Vout, H, const AXIS: usize>(
        aabb: Aabb<T>,
        aabb_center: [T; 3],
        input_vertices: &[impl Into<[T; 3]> + Clone + Sync],
        tolerance: T,
        used_vertices: Option<Vob>,
    ) -> Result<(Vec<Vout>, Vec<Index>), DeDupError>
    where
        T: ScalarKernel,
        Index: IndexType,
        Vout: Into<[T; 3]> + From<[T; 3]> + Clone + Sync,
        H: HashType<T>,
        usize: AsPrimitive<T>,
    {
        // totally arbitrary numbers
        let num_shards: usize = (rayon::current_num_threads() * 2).min(24);
        let min_axis = aabb.min[AXIS] - aabb_center[AXIS] - tolerance * 2.as_();
        let range = aabb.max[AXIS] - aabb_center[AXIS] - min_axis + tolerance * 4.as_();
        /*println!(
            "min_axis:{min_axis:?}, max_axis:{:?}, num_shards:{num_shards:?} Axis:{AXIS}",
            aabb.max[AXIS] - aabb_center[AXIS]
        );*/
        // Create shards with overlap
        let shard_width = range / num_shards.as_();

        let mut shards: Vec<_> = (0..num_shards)
            .map(|i| {
                let shard_min = min_axis + i.as_() * shard_width;
                let shard_max = min_axis + (i + 1).as_() * shard_width;
                //println!("shard_min:{shard_min:?}, shard_max:{shard_max:?}");
                Shard::<T, H, AXIS>::new(shard_min, shard_max)
            })
            .collect();

        /*println!(
            "first shard min:{:?}", &shards[0].min_axis,
        );*/
        // Expand the last shard to include the boundary
        if let Some(last_shard) = shards.last_mut() {
            last_shard.max_axis += tolerance;
            /*println!(
                "last shard max:{:?}", &last_shard.max_axis,
            );*/
        }

        // Process shards in parallel
        let all_clusters: Vec<_> = shards
            .par_iter()
            .map(|shard| {
                shard.deduplicate_shard(
                    tolerance,
                    input_vertices,
                    aabb_center,
                    &used_vertices,
                    num_shards,
                )
            })
            .flatten()
            .collect();

        // a flag indicated already re-assigned (original) indices/vertices
        let mut processed_original = Vob::from_elem(false, input_vertices.len());
        // a list of de-duplicated vertices
        let mut new_vertices: Vec<Vout> = Vec::with_capacity(input_vertices.len());
        // old vertex index -> new vertex index
        let mut new_map = vec![Index::MAX; input_vertices.len()];

        let max_centroid_len = all_clusters
            .iter()
            .map(|x| x.centroids.len())
            .max()
            .unwrap_or_default();
        // Local mapping: shard_centroid_index -> new_vertex_index
        let mut centroid_to_new_index = vec![Index::MAX; max_centroid_len];

        // collect the results
        for cluster in all_clusters {
            //println!("cluster.indices: {:?}", cluster.indices);
            //println!("cluster.remap: {:?}", cluster.remap);
            //println!("cluster.centroids: {:?}", cluster.centroids);

            let mut added_centroid = Vob::from_elem(false, cluster.centroids.len());

            // First pass: add unique centroids to new_vertices and build local mapping
            for (pos, vertex_id) in cluster.indices.iter().enumerate() {
                if !processed_original.ᚦget(*vertex_id) {
                    let centroid_index = cluster.remap[pos];
                    if !added_centroid.ᚦget(centroid_index) {
                        added_centroid.ᚦset(centroid_index, true);
                        centroid_to_new_index[centroid_index] =
                            Index::from_usize(new_vertices.len());
                        new_vertices.push(cluster.centroids[centroid_index].into());
                    }
                }
            }

            //println!("added_centroid:{added_centroid:?}");
            //println!("centroid_to_new_index:{centroid_to_new_index:?}");

            // Second pass: assign new indices to vertices using the local mapping
            for (pos, vertex_id) in cluster.indices.iter().enumerate() {
                if !processed_original.get(*vertex_id).unwrap() {
                    let _ = processed_original.set(*vertex_id, true);
                    let centroid_index = cluster.remap[pos];
                    new_map[*vertex_id] = centroid_to_new_index[centroid_index];
                }
            }
            //println!("new_map:{new_map:?}");
        }
        //println!("final new_map:{new_map:?}");
        Ok((new_vertices, new_map))
    }

    fn dedup_vertices_exact<T, Index, Vout, CheckFinite>(
        _vertices: &[impl Into<[T; 3]> + Clone + Sync],
        _unused_vertices: Option<Vob>,
    ) -> Result<(Vec<Vout>, Vec<Index>), DeDupError>
    where
        T: ScalarKernel,
        Index: IndexType,
        CheckFinite: CheckFinitePolicy,
    {
        Err(DeDupError("Not implemented".to_string()))
    }

    fn get_unused_vertices<T, Index>(
        vertices: &[impl Into<[T; 3]> + Clone + Sync],
        indices: &[Index],
    ) -> Result<Option<Vob>, DeDupError>
    where
        T: ScalarKernel,
        Index: IndexType,
    {
        SingleThreaded::get_unused_vertices::<T, Index>(vertices, indices)
    }
}