//! K-D Tree implementation for spatial queries.
//!
//! This module provides an efficient data structure for nearest-neighbor
//! searches in 3D space, used for mapping MRE properties to mesh nodes.

use nalgebra::{Point3, Vector3};
use serde::{Deserialize, Serialize};

/// A node in the K-D tree.
#[derive(Debug, Clone, Serialize, Deserialize)]
struct KdNode {
    /// The point stored at this node.
    point: Point3<f64>,
    /// Index of this point in the original data.
    index: usize,
    /// Left child (points with coordinate < split value).
    left: Option<Box<KdNode>>,
    /// Right child (points with coordinate >= split value).
    right: Option<Box<KdNode>>,
    /// Split dimension (0=x, 1=y, 2=z).
    split_dim: usize,
}

/// A 3D K-D Tree for efficient nearest-neighbor queries.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KdTree3D {
    root: Option<Box<KdNode>>,
    size: usize,
}

impl KdTree3D {
    /// Build a K-D tree from a set of points.
    pub fn build(points: &[Point3<f64>]) -> Self {
        if points.is_empty() {
            return Self {
                root: None,
                size: 0,
            };
        }

        let indices: Vec<usize> = (0..points.len()).collect();
        let root = Self::build_recursive(points, &indices, 0);

        Self {
            root: Some(root),
            size: points.len(),
        }
    }

    fn build_recursive(points: &[Point3<f64>], indices: &[usize], depth: usize) -> Box<KdNode> {
        let dim = depth % 3;

        // Sort indices by the current dimension
        let mut sorted_indices = indices.to_vec();
        sorted_indices.sort_by(|&a, &b| {
            let val_a = points[a].coords[dim];
            let val_b = points[b].coords[dim];
            val_a.partial_cmp(&val_b).unwrap()
        });

        let mid = sorted_indices.len() / 2;
        let median_idx = sorted_indices[mid];

        let left = if mid > 0 {
            Some(Self::build_recursive(
                points,
                &sorted_indices[..mid],
                depth + 1,
            ))
        } else {
            None
        };

        let right = if mid + 1 < sorted_indices.len() {
            Some(Self::build_recursive(
                points,
                &sorted_indices[mid + 1..],
                depth + 1,
            ))
        } else {
            None
        };

        Box::new(KdNode {
            point: points[median_idx],
            index: median_idx,
            left,
            right,
            split_dim: dim,
        })
    }

    /// Get the number of points in the tree.
    pub fn len(&self) -> usize {
        self.size
    }

    /// Check if the tree is empty.
    pub fn is_empty(&self) -> bool {
        self.size == 0
    }

    /// Find the nearest neighbor to a query point.
    ///
    /// Returns (index, distance) of the nearest point.
    pub fn nearest(&self, query: &Point3<f64>) -> Option<(usize, f64)> {
        let root = self.root.as_ref()?;
        let mut best = (root.index, (root.point - query).norm());
        Self::nearest_recursive(root, query, &mut best);
        Some(best)
    }

    fn nearest_recursive(node: &KdNode, query: &Point3<f64>, best: &mut (usize, f64)) {
        let dist = (node.point - query).norm();
        if dist < best.1 {
            *best = (node.index, dist);
        }

        let dim = node.split_dim;
        let diff = query.coords[dim] - node.point.coords[dim];

        // Determine which child to search first
        let (first, second) = if diff < 0.0 {
            (&node.left, &node.right)
        } else {
            (&node.right, &node.left)
        };

        // Search the closer subtree first
        if let Some(child) = first {
            Self::nearest_recursive(child, query, best);
        }

        // Only search the other subtree if it could contain a closer point
        if diff.abs() < best.1 {
            if let Some(child) = second {
                Self::nearest_recursive(child, query, best);
            }
        }
    }

    /// Find all points within a given radius of a query point.
    ///
    /// Returns a vector of (index, distance) pairs.
    pub fn within_radius(&self, query: &Point3<f64>, radius: f64) -> Vec<(usize, f64)> {
        let mut results = Vec::new();
        if let Some(root) = &self.root {
            Self::within_radius_recursive(root, query, radius, &mut results);
        }
        results
    }

    fn within_radius_recursive(
        node: &KdNode,
        query: &Point3<f64>,
        radius: f64,
        results: &mut Vec<(usize, f64)>,
    ) {
        let dist = (node.point - query).norm();
        if dist <= radius {
            results.push((node.index, dist));
        }

        let dim = node.split_dim;
        let diff = query.coords[dim] - node.point.coords[dim];

        // Check if we need to search left subtree
        if diff - radius <= 0.0 {
            if let Some(left) = &node.left {
                Self::within_radius_recursive(left, query, radius, results);
            }
        }

        // Check if we need to search right subtree
        if diff + radius >= 0.0 {
            if let Some(right) = &node.right {
                Self::within_radius_recursive(right, query, radius, results);
            }
        }
    }

    /// Find the k nearest neighbors to a query point.
    ///
    /// Returns a vector of (index, distance) pairs, sorted by distance.
    pub fn k_nearest(&self, query: &Point3<f64>, k: usize) -> Vec<(usize, f64)> {
        if k == 0 || self.is_empty() {
            return Vec::new();
        }

        let mut results: Vec<(usize, f64)> = Vec::with_capacity(k);
        if let Some(root) = &self.root {
            Self::k_nearest_recursive(root, query, k, &mut results);
        }

        results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
        results
    }

    fn k_nearest_recursive(
        node: &KdNode,
        query: &Point3<f64>,
        k: usize,
        results: &mut Vec<(usize, f64)>,
    ) {
        let dist = (node.point - query).norm();

        // Check if we should add this point
        if results.len() < k {
            results.push((node.index, dist));
            results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
        } else if dist < results.last().unwrap().1 {
            results.pop();
            results.push((node.index, dist));
            results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
        }

        let dim = node.split_dim;
        let diff = query.coords[dim] - node.point.coords[dim];
        let max_dist = if results.len() < k {
            f64::INFINITY
        } else {
            results.last().unwrap().1
        };

        // Determine which child to search first
        let (first, second) = if diff < 0.0 {
            (&node.left, &node.right)
        } else {
            (&node.right, &node.left)
        };

        // Search the closer subtree first
        if let Some(child) = first {
            Self::k_nearest_recursive(child, query, k, results);
        }

        // Update max distance after searching first subtree
        let max_dist = if results.len() < k {
            f64::INFINITY
        } else {
            results.last().unwrap().1
        };

        // Only search the other subtree if it could contain a closer point
        if diff.abs() < max_dist {
            if let Some(child) = second {
                Self::k_nearest_recursive(child, query, k, results);
            }
        }
    }
}

/// Compute the bounding box of a set of points.
pub fn bounding_box(points: &[Point3<f64>]) -> Option<(Point3<f64>, Point3<f64>)> {
    if points.is_empty() {
        return None;
    }

    let mut min = points[0];
    let mut max = points[0];

    for p in points.iter().skip(1) {
        for i in 0..3 {
            if p.coords[i] < min.coords[i] {
                min.coords[i] = p.coords[i];
            }
            if p.coords[i] > max.coords[i] {
                max.coords[i] = p.coords[i];
            }
        }
    }

    Some((min, max))
}

/// Compute the center of mass of a set of points.
pub fn center_of_mass(points: &[Point3<f64>]) -> Option<Point3<f64>> {
    if points.is_empty() {
        return None;
    }

    let sum: Vector3<f64> = points.iter().map(|p| p.coords).sum();
    Some(Point3::from(sum / points.len() as f64))
}

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

    fn create_test_points() -> Vec<Point3<f64>> {
        vec![
            Point3::new(0.0, 0.0, 0.0),
            Point3::new(1.0, 0.0, 0.0),
            Point3::new(0.0, 1.0, 0.0),
            Point3::new(0.0, 0.0, 1.0),
            Point3::new(1.0, 1.0, 0.0),
            Point3::new(1.0, 0.0, 1.0),
            Point3::new(0.0, 1.0, 1.0),
            Point3::new(1.0, 1.0, 1.0),
        ]
    }

    #[test]
    fn test_build_tree() {
        let points = create_test_points();
        let tree = KdTree3D::build(&points);

        assert_eq!(tree.len(), 8);
        assert!(!tree.is_empty());
    }

    #[test]
    fn test_empty_tree() {
        let tree = KdTree3D::build(&[]);
        assert!(tree.is_empty());
        assert_eq!(tree.len(), 0);
        assert!(tree.nearest(&Point3::new(0.0, 0.0, 0.0)).is_none());
    }

    #[test]
    fn test_nearest_exact_match() {
        let points = create_test_points();
        let tree = KdTree3D::build(&points);

        // Query for a point that exists in the tree
        let result = tree.nearest(&Point3::new(1.0, 1.0, 1.0)).unwrap();
        assert_eq!(result.0, 7); // Index of (1,1,1)
        assert!(result.1 < 1e-10);
    }

    #[test]
    fn test_nearest_interpolated() {
        let points = create_test_points();
        let tree = KdTree3D::build(&points);

        // Query for a point not in the tree
        let result = tree.nearest(&Point3::new(0.1, 0.1, 0.1)).unwrap();
        assert_eq!(result.0, 0); // Closest to origin
    }

    #[test]
    fn test_within_radius() {
        let points = create_test_points();
        let tree = KdTree3D::build(&points);

        // All points at distance <= sqrt(3) from origin
        let results = tree.within_radius(&Point3::new(0.0, 0.0, 0.0), 2.0);
        assert_eq!(results.len(), 8);

        // Only origin at distance 0
        let results = tree.within_radius(&Point3::new(0.0, 0.0, 0.0), 0.1);
        assert_eq!(results.len(), 1);
    }

    #[test]
    fn test_k_nearest() {
        let points = create_test_points();
        let tree = KdTree3D::build(&points);

        // Find 3 nearest to origin
        let results = tree.k_nearest(&Point3::new(0.0, 0.0, 0.0), 3);
        assert_eq!(results.len(), 3);

        // First should be origin
        assert_eq!(results[0].0, 0);
        assert!(results[0].1 < 1e-10);

        // Results should be sorted by distance
        for i in 1..results.len() {
            assert!(results[i].1 >= results[i - 1].1);
        }
    }

    #[test]
    fn test_bounding_box() {
        let points = create_test_points();
        let (min, max) = bounding_box(&points).unwrap();

        assert!((min - Point3::new(0.0, 0.0, 0.0)).norm() < 1e-10);
        assert!((max - Point3::new(1.0, 1.0, 1.0)).norm() < 1e-10);
    }

    #[test]
    fn test_center_of_mass() {
        let points = create_test_points();
        let com = center_of_mass(&points).unwrap();

        assert!((com - Point3::new(0.5, 0.5, 0.5)).norm() < 1e-10);
    }
}