threecrate-algorithms 0.10.0

//! Normal estimation algorithms

use crate::nearest_neighbor::KdTree;
use nalgebra::Matrix3;
use rayon::prelude::*;
use threecrate_core::{
    Error, NearestNeighborSearch, NormalPoint3f, Point3f, PointCloud, Result, Vector3f,
};

#[cfg(test)]
use std::cmp::Ordering;
#[cfg(test)]
use std::collections::BinaryHeap;

/// Configuration for normal estimation
#[derive(Debug, Clone)]
pub struct NormalEstimationConfig {
    /// Number of nearest neighbors to use (k-NN)
    pub k_neighbors: usize,
    /// Optional radius for radius-based neighbor search
    pub radius: Option<f32>,
    /// Whether to enforce orientation consistency
    pub consistent_orientation: bool,
    /// Viewpoint for orientation consistency (if None, uses positive Z direction)
    pub viewpoint: Option<Point3f>,
}

impl Default for NormalEstimationConfig {
    fn default() -> Self {
        Self {
            k_neighbors: 10,
            radius: None,
            consistent_orientation: true,
            viewpoint: None,
        }
    }
}

/// A simple distance-based neighbor for priority queue
#[derive(Debug, Clone)]
#[cfg(test)]
struct Neighbor {
    index: usize,
    distance: f32,
}

#[cfg(test)]
impl PartialEq for Neighbor {
    fn eq(&self, other: &Self) -> bool {
        self.distance == other.distance
    }
}

#[cfg(test)]
impl Eq for Neighbor {}

#[cfg(test)]
impl PartialOrd for Neighbor {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[cfg(test)]
impl Ord for Neighbor {
    fn cmp(&self, other: &Self) -> Ordering {
        // Reverse ordering for min-heap behavior
        other
            .distance
            .partial_cmp(&self.distance)
            .unwrap_or(Ordering::Equal)
    }
}

/// Find k-nearest neighbors using brute force search
#[cfg(test)]
fn find_k_nearest_neighbors(points: &[Point3f], query_idx: usize, k: usize) -> Vec<usize> {
    let query = &points[query_idx];
    let mut heap = BinaryHeap::with_capacity(k + 1);

    for (i, point) in points.iter().enumerate() {
        if i == query_idx {
            continue; // Skip the query point itself
        }

        let distance = (point - query).magnitude_squared();
        let neighbor = Neighbor { index: i, distance };

        if heap.len() < k {
            heap.push(neighbor);
        } else if let Some(farthest) = heap.peek() {
            if neighbor.distance < farthest.distance {
                heap.pop();
                heap.push(neighbor);
            }
        }
    }

    heap.into_iter().map(|n| n.index).collect()
}

/// Find neighbors within a radius
#[cfg(test)]
fn find_radius_neighbors(points: &[Point3f], query_idx: usize, radius: f32) -> Vec<usize> {
    let query = &points[query_idx];
    let radius_squared = radius * radius;

    points
        .iter()
        .enumerate()
        .filter(|(i, point)| {
            *i != query_idx && (**point - query).magnitude_squared() <= radius_squared
        })
        .map(|(i, _)| i)
        .collect()
}

/// Find neighbors using either k-NN or radius-based search
#[cfg(test)]
fn find_neighbors(
    points: &[Point3f],
    query_idx: usize,
    config: &NormalEstimationConfig,
) -> Vec<usize> {
    if let Some(radius) = config.radius {
        // Use radius-based search
        find_radius_neighbors(points, query_idx, radius)
    } else {
        // Use k-NN search
        find_k_nearest_neighbors(points, query_idx, config.k_neighbors)
    }
}

/// Find neighbors using a prebuilt KD-tree.
/// Fill `neighborhood` with the `k` nearest neighbors of point `index`, not
/// counting the point itself. `knn` is scratch space, reused across calls so the
/// search does not allocate per point.
fn nearest_excluding(
    tree: &KdTree,
    index: usize,
    query: &Point3f,
    k: usize,
    knn: &mut Vec<(usize, f32)>,
    neighborhood: &mut Vec<usize>,
) {
    tree.find_k_nearest_into(query, k + 1, knn);
    neighborhood.clear();
    neighborhood.extend(
        knn.iter()
            .map(|&(idx, _)| idx)
            .filter(|&idx| idx != index)
            .take(k),
    );
}

/// Compute normal using PCA on the neighborhood points
fn compute_normal_pca(points: &[Point3f], indices: &[usize]) -> Vector3f {
    if indices.len() < 3 {
        // Default normal if not enough points
        return Vector3f::new(0.0, 0.0, 1.0);
    }

    // Compute centroid
    let mut centroid = Point3f::origin();
    for &idx in indices {
        centroid += points[idx].coords;
    }
    centroid /= indices.len() as f32;

    // Build covariance matrix
    let mut covariance = Matrix3::zeros();
    for &idx in indices {
        let diff = points[idx] - centroid;
        covariance += diff * diff.transpose();
    }
    covariance /= indices.len() as f32;

    // Find eigenvector corresponding to smallest eigenvalue
    // This is the normal direction
    let eigen = covariance.symmetric_eigen();
    let eigenvalues = eigen.eigenvalues;
    let eigenvectors = eigen.eigenvectors;

    // Find index of smallest eigenvalue
    let mut min_idx = 0;
    for i in 1..3 {
        if eigenvalues[i] < eigenvalues[min_idx] {
            min_idx = i;
        }
    }

    // Return the corresponding eigenvector as normal
    let mut normal: Vector3f = eigenvectors.column(min_idx).into();

    // Ensure the normal is normalized
    let magnitude = normal.magnitude();
    if magnitude > 1e-6 {
        normal /= magnitude;
    } else {
        normal = Vector3f::new(0.0, 0.0, 1.0);
    }

    normal
}

/// Orient normal towards viewpoint for consistency
fn orient_normal_towards_viewpoint(
    normal: Vector3f,
    point: Point3f,
    viewpoint: Point3f,
) -> Vector3f {
    let to_viewpoint = (viewpoint - point).normalize();
    let dot_product = normal.dot(&to_viewpoint);

    // If the angle between normal and viewpoint direction is > 90 degrees, flip the normal
    if dot_product < 0.0 {
        -normal
    } else {
        normal
    }
}

/// Estimate normals for a point cloud using k-nearest neighbors
///
/// This function computes surface normals for each point in the cloud by:
/// 1. Finding k nearest neighbors for each point (or neighbors within radius)
/// 2. Computing the normal using Principal Component Analysis (PCA)
/// 3. The normal is the eigenvector corresponding to the smallest eigenvalue
/// 4. Optionally enforcing orientation consistency towards a viewpoint
///
/// # Arguments
/// * `cloud` - Reference to the point cloud
/// * `k` - Number of nearest neighbors to use (typically 10-30)
///
/// # Returns
/// * `Result<PointCloud<NormalPoint3f>>` - A new point cloud with normals
pub fn estimate_normals(
    cloud: &PointCloud<Point3f>,
    k: usize,
) -> Result<PointCloud<NormalPoint3f>> {
    let config = NormalEstimationConfig {
        k_neighbors: k,
        ..Default::default()
    };
    estimate_normals_with_config(cloud, &config)
}

/// Estimate normals with advanced configuration
///
/// # Arguments
/// * `cloud` - Reference to the point cloud
/// * `config` - Configuration for normal estimation
///
/// # Returns
/// * `Result<PointCloud<NormalPoint3f>>` - A new point cloud with normals
pub fn estimate_normals_with_config(
    cloud: &PointCloud<Point3f>,
    config: &NormalEstimationConfig,
) -> Result<PointCloud<NormalPoint3f>> {
    if cloud.is_empty() {
        return Ok(PointCloud::new());
    }

    if config.k_neighbors < 3 {
        return Err(Error::InvalidData(
            "k_neighbors must be at least 3".to_string(),
        ));
    }

    let points = &cloud.points;
    let search_tree = KdTree::new(points)?;

    // Determine viewpoint for orientation consistency
    let viewpoint = config.viewpoint.unwrap_or_else(|| {
        // Default viewpoint: compute a good viewpoint based on the point cloud bounds
        let mut min_x = points[0].x;
        let mut min_y = points[0].y;
        let mut min_z = points[0].z;
        let mut max_x = points[0].x;
        let mut max_y = points[0].y;
        let mut max_z = points[0].z;

        for point in points {
            min_x = min_x.min(point.x);
            min_y = min_y.min(point.y);
            min_z = min_z.min(point.z);
            max_x = max_x.max(point.x);
            max_y = max_y.max(point.y);
            max_z = max_z.max(point.z);
        }

        let center = Point3f::new(
            (min_x + max_x) / 2.0,
            (min_y + max_y) / 2.0,
            (min_z + max_z) / 2.0,
        );
        let extent =
            ((max_x - min_x).powi(2) + (max_y - min_y).powi(2) + (max_z - min_z).powi(2)).sqrt();

        // Viewpoint is above the center of the point cloud
        center + Vector3f::new(0.0, 0.0, extent)
    });

    // Compute normals in parallel. Each rayon thread keeps its own scratch
    // buffers, so the per-point neighbor search does not allocate.
    let normals: Vec<NormalPoint3f> = (0..points.len())
        .into_par_iter()
        .map_init(
            || (Vec::new(), Vec::new()),
            |(knn, neighborhood): &mut (Vec<(usize, f32)>, Vec<usize>), i| {
                neighborhood.clear();
                if let Some(radius) = config.radius {
                    neighborhood.extend(
                        search_tree
                            .find_radius_neighbors(&points[i], radius)
                            .into_iter()
                            .map(|(idx, _)| idx)
                            .filter(|&idx| idx != i),
                    );
                }

                // k-NN: the default, and the fallback when the radius finds too few
                if config.radius.is_none() || neighborhood.len() < config.k_neighbors {
                    nearest_excluding(
                        &search_tree,
                        i,
                        &points[i],
                        config.k_neighbors,
                        knn,
                        neighborhood,
                    );
                }

                // Ensure we have enough neighbors for PCA
                if neighborhood.len() < 3 {
                    let fallback_k = config.k_neighbors.max(5);
                    nearest_excluding(&search_tree, i, &points[i], fallback_k, knn, neighborhood);
                }

                // Include the point itself in its neighborhood
                neighborhood.push(i);

                let mut normal = compute_normal_pca(points, neighborhood);

                // Apply orientation consistency if requested
                if config.consistent_orientation {
                    normal = orient_normal_towards_viewpoint(normal, points[i], viewpoint);
                }

                NormalPoint3f {
                    position: points[i],
                    normal,
                }
            },
        )
        .collect();

    Ok(PointCloud::from_points(normals))
}

/// Flip normals so neighbouring normals agree, all the way around the object.
///
/// `estimate_normals` points every normal toward one viewpoint, which is right
/// for a single scan but leaves about half the normals of a closed object
/// pointing inward. This walks a minimum spanning tree over the `k` nearest
/// neighbours, preferring nearly parallel pairs, and flips each normal to agree
/// with the one it was reached from (Hoppe et al. 1992). Each connected piece
/// starts from its highest point, whose normal is turned upward, so closed
/// objects end up with outward normals. Positions are not changed.
pub fn orient_normals_consistent(cloud: &mut PointCloud<NormalPoint3f>, k: usize) -> Result<()> {
    use std::cmp::Reverse;
    use std::collections::BinaryHeap;

    let n = cloud.points.len();
    if n < 2 {
        return Ok(());
    }
    if k < 1 {
        return Err(Error::InvalidData("k must be at least 1".to_string()));
    }
    if n > u32::MAX as usize {
        return Err(Error::InvalidData(
            "too many points to orient (limit is about 4 billion)".to_string(),
        ));
    }
    let positions: Vec<Point3f> = cloud.points.iter().map(|p| p.position).collect();
    let tree = KdTree::new(&positions)?;
    let k = k.min(n - 1);
    // One flat row of `k` neighbours per point (u32 to halve memory on large
    // clouds); NONE pads rows with fewer neighbours.
    const NONE: u32 = u32::MAX;
    let mut neighbours = vec![NONE; n * k];
    neighbours
        .par_chunks_mut(k)
        .enumerate()
        .for_each_init(Vec::new, |found, (i, row)| {
            tree.find_k_nearest_into(&positions[i], k + 1, found);
            let others = found.iter().map(|&(j, _)| j).filter(|&j| j != i);
            for (slot, j) in row.iter_mut().zip(others) {
                *slot = j as u32;
            }
        });

    /// A candidate tree edge (cost, from, to); cheaper edges join more
    /// parallel normals.
    #[derive(PartialEq)]
    struct Edge(f32, u32, u32);
    impl Eq for Edge {}
    impl PartialOrd for Edge {
        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
            Some(self.cmp(other))
        }
    }
    impl Ord for Edge {
        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
            self.0.total_cmp(&other.0)
        }
    }

    let points = &mut cloud.points;
    let mut visited = vec![false; n];
    let mut heap = BinaryHeap::new();
    // Queue edges only towards points not yet in the tree, so the heap stays
    // far smaller than n * k.
    let push = |heap: &mut BinaryHeap<Reverse<Edge>>,
                from: usize,
                points: &[NormalPoint3f],
                visited: &[bool]| {
        for &to in &neighbours[from * k..(from + 1) * k] {
            if to != NONE && !visited[to as usize] {
                let cost = 1.0 - points[from].normal.dot(&points[to as usize].normal).abs();
                heap.push(Reverse(Edge(cost, from as u32, to)));
            }
        }
    };
    // Seeds, highest first, so every piece starts from its top.
    let mut seeds: Vec<usize> = (0..n).collect();
    seeds.sort_unstable_by(|&a, &b| positions[b].z.total_cmp(&positions[a].z));

    for seed in seeds {
        if visited[seed] {
            continue;
        }
        visited[seed] = true;
        if points[seed].normal.z < 0.0 {
            points[seed].normal = -points[seed].normal;
        }
        push(&mut heap, seed, points, &visited);
        while let Some(Reverse(Edge(_, from, to))) = heap.pop() {
            let (from, to) = (from as usize, to as usize);
            if visited[to] {
                continue;
            }
            visited[to] = true;
            if points[from].normal.dot(&points[to].normal) < 0.0 {
                points[to].normal = -points[to].normal;
            }
            push(&mut heap, to, points, &visited);
        }
    }
    Ok(())
}

/// Estimate normals using radius-based neighbor search
///
/// # Arguments
/// * `cloud` - Reference to the point cloud
/// * `radius` - Search radius for neighbors
/// * `consistent_orientation` - Whether to enforce orientation consistency
///
/// # Returns
/// * `Result<PointCloud<NormalPoint3f>>` - A new point cloud with normals
pub fn estimate_normals_radius(
    cloud: &PointCloud<Point3f>,
    radius: f32,
    consistent_orientation: bool,
) -> Result<PointCloud<NormalPoint3f>> {
    let config = NormalEstimationConfig {
        k_neighbors: 10, // Fallback value
        radius: Some(radius),
        consistent_orientation,
        viewpoint: None,
    };
    estimate_normals_with_config(cloud, &config)
}

/// Estimate normals and modify the input cloud in-place (legacy API)
/// This function is deprecated in favor of the version that returns a new cloud
#[deprecated(note = "Use estimate_normals instead which returns a new point cloud")]
pub fn estimate_normals_inplace(_cloud: &mut PointCloud<Point3f>, k: usize) -> Result<()> {
    // This would require converting the point cloud type, which isn't straightforward
    // with the current type system. The new API is cleaner.
    let _ = k;
    Err(Error::Unsupported(
        "Use estimate_normals instead".to_string(),
    ))
}

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

    #[test]
    fn test_estimate_normals_simple() {
        // Create a simple planar point cloud (XY plane)
        let mut cloud = PointCloud::new();
        cloud.push(Point3f::new(0.0, 0.0, 0.0));
        cloud.push(Point3f::new(1.0, 0.0, 0.0));
        cloud.push(Point3f::new(0.0, 1.0, 0.0));
        cloud.push(Point3f::new(1.0, 1.0, 0.0));
        cloud.push(Point3f::new(0.5, 0.5, 0.0));

        let result = estimate_normals(&cloud, 3).unwrap();

        assert_eq!(result.len(), 5);

        // For a planar surface in XY plane, normals should point along Z axis
        for point in result.iter() {
            let normal = point.normal;
            // Normal should be close to (0, 0, 1) or (0, 0, -1)
            assert!(
                normal.z.abs() > 0.8,
                "Normal should be primarily in Z direction: {:?}",
                normal
            );
        }
    }

    #[test]
    fn test_estimate_normals_empty() {
        let cloud = PointCloud::<Point3f>::new();
        let result = estimate_normals(&cloud, 5).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_estimate_normals_insufficient_k() {
        let mut cloud = PointCloud::new();
        cloud.push(Point3f::new(0.0, 0.0, 0.0));

        let result = estimate_normals(&cloud, 2);
        assert!(result.is_err());
    }

    #[test]
    fn test_estimate_normals_radius() {
        // Create a simple planar point cloud for testing radius-based search
        let mut cloud = PointCloud::new();
        for i in 0..20 {
            for j in 0..20 {
                let x = (i as f32) * 0.1;
                let y = (j as f32) * 0.1;
                let z = 0.0;
                cloud.push(Point3f::new(x, y, z));
            }
        }

        let result = estimate_normals_radius(&cloud, 0.2, true).unwrap();
        assert_eq!(result.len(), 400);

        // Check that normals are computed and have reasonable values
        let mut z_direction_count = 0;
        for point in result.iter() {
            let normal_magnitude = point.normal.magnitude();
            // Normals should be unit vectors
            assert!(
                (normal_magnitude - 1.0).abs() < 0.1,
                "Normal should be unit vector: magnitude={}",
                normal_magnitude
            );

            // For a planar surface, normals should be primarily in Z direction
            if point.normal.z.abs() > 0.8 {
                z_direction_count += 1;
            }
        }

        // At least 80% of normals should be in Z direction for a planar surface
        let percentage = (z_direction_count as f32 / result.len() as f32) * 100.0;
        assert!(
            percentage > 80.0,
            "Only {:.1}% of normals are in Z direction",
            percentage
        );
    }

    #[test]
    fn test_estimate_normals_cylinder() {
        // Create a true cylindrical point cloud (points on a cylinder surface)
        let mut cloud = PointCloud::new();
        for i in 0..10 {
            for j in 0..10 {
                let angle = (i as f32) * 0.6;
                let height = (j as f32) * 0.2 - 1.0;
                let x = angle.cos();
                let y = angle.sin();
                let z = height;
                cloud.push(Point3f::new(x, y, z));
            }
        }

        let config = NormalEstimationConfig {
            k_neighbors: 8, // Increase k for better results
            radius: None,
            consistent_orientation: true,
            viewpoint: Some(Point3f::new(0.0, 0.0, 2.0)), // View from above
        };

        let result = estimate_normals_with_config(&cloud, &config).unwrap();
        assert_eq!(result.len(), 100);

        // Check that normals are computed and have reasonable values
        let mut perpendicular_count = 0;
        let mut outward_count = 0;
        for point in result.iter() {
            let normal_magnitude = point.normal.magnitude();
            // Normals should be unit vectors
            assert!(
                (normal_magnitude - 1.0).abs() < 0.1,
                "Normal should be unit vector: magnitude={}",
                normal_magnitude
            );

            // For a cylinder, normals should be roughly perpendicular to the cylinder axis (Z-axis)
            let dot_with_z = point.normal.z.abs();
            if dot_with_z < 0.8 {
                perpendicular_count += 1;
            }

            // Check if normal points outward from center
            let to_center = Vector3f::new(-point.position.x, -point.position.y, 0.0).normalize();
            let dot_outward = point.normal.dot(&to_center);
            if dot_outward > 0.5 {
                outward_count += 1;
            }
        }

        // At least 60% of normals should be perpendicular to Z-axis for a cylinder
        let percentage_perpendicular = (perpendicular_count as f32 / result.len() as f32) * 100.0;
        let percentage_outward = (outward_count as f32 / result.len() as f32) * 100.0;

        println!(
            "Cylinder test: {:.1}% perpendicular to Z, {:.1}% pointing outward",
            percentage_perpendicular, percentage_outward
        );

        // For a cylinder, normals should be perpendicular to Z-axis
        assert!(
            percentage_perpendicular > 60.0,
            "Only {:.1}% of normals are perpendicular to Z-axis",
            percentage_perpendicular
        );
    }

    #[test]
    fn test_estimate_normals_orientation_consistency() {
        // Create a simple planar point cloud
        let mut cloud = PointCloud::new();
        cloud.push(Point3f::new(0.0, 0.0, 0.0));
        cloud.push(Point3f::new(1.0, 0.0, 0.0));
        cloud.push(Point3f::new(0.0, 1.0, 0.0));
        cloud.push(Point3f::new(1.0, 1.0, 0.0));

        // Test with orientation consistency enabled
        let config_consistent = NormalEstimationConfig {
            k_neighbors: 3,
            radius: None,
            consistent_orientation: true,
            viewpoint: Some(Point3f::new(0.0, 0.0, 1.0)), // View from positive Z
        };

        let result_consistent = estimate_normals_with_config(&cloud, &config_consistent).unwrap();

        // Test with orientation consistency disabled
        let config_inconsistent = NormalEstimationConfig {
            k_neighbors: 3,
            radius: None,
            consistent_orientation: false,
            viewpoint: None,
        };

        let _result_inconsistent =
            estimate_normals_with_config(&cloud, &config_inconsistent).unwrap();

        // With consistent orientation, all normals should point in the same direction (positive Z)
        let first_normal_consistent = result_consistent.points[0].normal.z;
        for point in result_consistent.iter() {
            assert!(
                (point.normal.z * first_normal_consistent) > 0.0,
                "Normals should have consistent orientation"
            );
        }

        // Without consistent orientation, normals might point in different directions
        // (This test is less strict since the algorithm might still produce consistent results)
        println!("Consistent orientation test completed");
    }

    #[test]
    fn test_find_neighbors() {
        let points = vec![
            Point3f::new(0.0, 0.0, 0.0),
            Point3f::new(1.0, 0.0, 0.0),
            Point3f::new(0.0, 1.0, 0.0),
            Point3f::new(2.0, 0.0, 0.0),
        ];

        // Test k-NN
        let config_knn = NormalEstimationConfig {
            k_neighbors: 2,
            radius: None,
            consistent_orientation: false,
            viewpoint: None,
        };

        let neighbors_knn = find_neighbors(&points, 0, &config_knn);
        assert_eq!(neighbors_knn.len(), 2);

        // Test radius-based
        let config_radius = NormalEstimationConfig {
            k_neighbors: 10,
            radius: Some(1.5),
            consistent_orientation: false,
            viewpoint: None,
        };

        let neighbors_radius = find_neighbors(&points, 0, &config_radius);
        assert_eq!(neighbors_radius.len(), 2); // Points at (1,0,0) and (0,1,0) are within radius 1.5
    }

    #[test]
    fn test_orient_normals_consistent_points_sphere_normals_outward() {
        let n = 2000;
        let golden = std::f32::consts::PI * (1.0 + 5f32.sqrt());
        let points: Vec<Point3f> = (0..n)
            .map(|i| {
                let t = i as f32 + 0.5;
                let phi = (1.0 - 2.0 * t / n as f32).acos();
                let theta = golden * t;
                Point3f::new(theta.cos() * phi.sin(), theta.sin() * phi.sin(), phi.cos())
            })
            .collect();
        let mut cloud = estimate_normals(&PointCloud::from_points(points), 10).unwrap();
        let outward = |c: &PointCloud<NormalPoint3f>| {
            c.points
                .iter()
                .filter(|p| p.normal.dot(&p.position.coords) > 0.0)
                .count()
        };
        // The viewpoint rule gets the lower half of the sphere wrong.
        assert!(outward(&cloud) < n * 9 / 10);
        orient_normals_consistent(&mut cloud, 10).unwrap();
        assert_eq!(outward(&cloud), n);
    }

    #[test]
    fn test_orient_normals_consistent_handles_tiny_clouds() {
        let mut empty = PointCloud::<NormalPoint3f>::new();
        orient_normals_consistent(&mut empty, 10).unwrap();
        let mut one = PointCloud::from_points(vec![NormalPoint3f {
            position: Point3f::origin(),
            normal: Vector3f::new(0.0, 0.0, -1.0),
        }]);
        orient_normals_consistent(&mut one, 10).unwrap();
    }
}