rust-3d 0.22.0

/*
Copyright 2016 Martin Buck
This file is part of rust-3d.
rust-3d is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
rust-3d is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.
You should have received a copy of the GNU Lesser General Public License
along with rust-3d.  If not, see <http://www.gnu.org/licenses/>.
*/

//! KdTree https://en.wikipedia.org/wiki/K-d_tree

use std::cmp::Ordering;

use prelude::*;
use distances_3d::*;
use functions::{dimension_compare, dimension_dist, sort_and_limit};

#[derive (Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
/// KdTree https://en.wikipedia.org/wiki/K-d_tree
pub struct KdTree<P> where
    P: Is3D {

    root: Option<KdNode<P>>
}

#[derive (Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct KdNode<P> where
    P: Is3D {

    pub left: Option<Box<KdNode<P>>>,
    pub right: Option<Box<KdNode<P>>>,
    pub val: P,
    pub dimension: i8
}

impl<P> IsTree3D<P> for KdTree<P> where
    P: Is3D + Clone {

    fn size(&self) -> usize {
        match self.root {
            None => 0,
            Some(ref node) => node.size()
        }
    }

    fn to_pointcloud(&self) -> PointCloud3D<P>{
        let mut result = PointCloud3D::new();
        if let Some(ref node) = self.root {
            node.to_pointcloud_3d(&mut result);
        }
        result
    }

    fn build(&mut self, pc: PointCloud3D<P>) -> Result<()> {
        match pc.len() {
            0 => Err(ErrorKind::TooFewPoints),
            _ => {
                self.root = Some(KdNode::new(0, pc.data));
                Ok(())
            }
        }
    }
}

impl<PSearch, PFind> IsKNearestSearchable<PSearch, PFind> for KdTree<PFind> where
    PSearch: Is3D,
    PFind:   Is3D + Clone {

    fn knearest(&self, search: &PSearch, n: usize) -> Vec<PFind> {
        let mut result = Vec::new();
        if n < 1 { return result; }
        if let Some(ref node) = self.root {
            node.knearest(search, n, &mut result);
        }
        return result;
    }

    fn nearest(&self, search: &PSearch) -> Result<PFind> { //@todo implemented on its own, since the code can be faster without vecs
        let result = self.knearest(search, 1);
        match result.len() {
            0 => Err(ErrorKind::TooFewPoints),
            _ => {
                let p = result[0].clone();
                Ok(p)
            }
        }
    }
}

impl<PSearch, PFind> IsSphereSearchable<PSearch, PFind> for KdTree<PFind> where
    PSearch: Is3D,
    PFind:   Is3D + Clone {

    fn in_sphere(&self, search: &PSearch, radius: f64) -> Vec<PFind> {
        let mut result = Vec::new();
        if radius <= 0.0 { return result; }
        if let Some(ref node) = self.root {
            node.in_sphere(search, radius, &mut result);
        }
        return result;
    }
}

impl<PSearch, PFind> IsBox3DSearchable<PSearch, PFind> for KdTree<PFind> where
    PSearch: Is3D,
    PFind:   Is3D + Clone {

    fn in_box(&self, search: &PSearch, x_size: f64, y_size: f64, z_size: f64) -> Vec<PFind> {
        let mut result = Vec::new();
        if x_size <= 0.0 || y_size <= 0.0 || z_size <= 0.0 { return result; }
        if let Some(ref node) = self.root {
            node.in_box(search, x_size, y_size, z_size, &mut result);
        }
        return result;
    }
}

impl<P> KdNode<P> where
    P: Is3D {

    pub fn size(&self) -> usize {
        let mut result: usize = 0;
        if let Some(ref n) = (&self).left { result += n.size(); }
        result += 1;
        if let Some(ref n) = (&self).right { result += n.size(); }
        result
    }

    fn is_leaf(&self) -> bool {
        self.left.is_none() && self.right.is_none()
    }
}

impl<P> KdNode<P> where
    P: Is3D + Clone {
    pub fn to_pointcloud_3d(&self, pc: &mut PointCloud3D<P>) {
        if let Some(ref n) = (&self).left { n.to_pointcloud_3d(pc); }
        pc.push(self.val.clone());
        if let Some(ref n) = (&self).right { n.to_pointcloud_3d(pc); }
    }

    pub fn new(dim: i8, mut pc: Vec<Box<P>>) -> KdNode<P> {
        let dimension = dim % 2;
        if pc.len() == 1 {
            return KdNode {
                left: None,
                right: None,
                val: *pc[0].clone(),
                dimension: dimension
            }
        }

        pc.sort_by(|a, b| match dimension {
            0 => a.x().partial_cmp(&b.x()).unwrap_or(Ordering::Equal),
            1 => a.y().partial_cmp(&b.y()).unwrap_or(Ordering::Equal),
            2 => a.z().partial_cmp(&b.z()).unwrap_or(Ordering::Equal),
            _ => Ordering::Equal
        });
        let median = pc.len() / 2;
        let mut pc_left = Vec::new();
        let mut pc_right = Vec::new();

        let val = *pc[median].clone();

        for (i, p) in pc.into_iter().enumerate() {
            if i < median { pc_left.push(p); } else if i > median { pc_right.push(p); }
        }

        let left = match pc_left.len() {
            0 => None,
            _ => Some(Box::new(KdNode::new(dimension + 1, pc_left)))
        };

        let right = match pc_right.len() {
            0 => None,
            _ => Some(Box::new(KdNode::new(dimension + 1, pc_right)))
        };

        KdNode {
            left: left,
            right: right,
            val: val,
            dimension: dimension
        }
    }
}

impl<PFind> KdNode<PFind> where
    PFind: Is3D + Clone {

    pub fn knearest<PSearch>(&self, search: &PSearch, n: usize, pc: &mut Vec<PFind>) where
        PSearch: Is3D {

        if pc.len() < n || sqr_dist_3d(search, &self.val) < sqr_dist_3d(search, &pc[&pc.len() -1 ]) {
            pc.push(self.val.clone());
        }

        let comp = dimension_compare(search, &self.val, self.dimension);

        match comp {
            Ok(res) => match res {
                Ordering::Less  => if let Some(ref node) = (&self).left { node.knearest(search, n, pc); },
                _               => if let Some(ref node) = (&self).right { node.knearest(search, n, pc); }
            },
            Err(_) => {}
        }

        sort_and_limit(pc, search, n);

        let (current_search, current_val) = match self.dimension {
            0 => (search.x(), self.val.x()),
            1 => (search.y(), self.val.y()),
            _ => (search.z(), self.val.z())
        };

        let distance_best = dist_3d(search, &pc[&pc.len() -1 ]);
        let border_left = current_search - distance_best;
        let border_right = current_search + distance_best;

        match comp {
            Ok(res) => match res {
                Ordering::Less => if let Some(ref node) = (&self).right {
                    if pc.len() < n || border_right >= current_val {
                        node.knearest(search, n, pc);
                    }
                },
                Ordering::Greater => if let Some(ref node) = (&self).left {
                    if pc.len() < n || border_left <= current_val {
                        node.knearest(search, n, pc);
                    }
                },
                Ordering::Equal => {}
            },
            Err(_) => {}
        }

        sort_and_limit(pc, search, n);
    }

    pub fn in_sphere<PSearch>(&self, search: &PSearch, radius: f64, pc: &mut Vec<PFind>) where
        PSearch: Is3D {

        if radius <= 0.0 { return; }

        if dist_3d(search, &self.val) <= radius {
            pc.push(self.val.clone());
        }

        if self.is_leaf() { return; }

        let comp = dimension_compare(search, &self.val, self.dimension);

        match comp {
            Ok(res) => match res {
                Ordering::Less  => if let Some(ref node) = (&self).left { node.in_sphere(search, radius, pc); },
                _               => if let Some(ref node) = (&self).right { node.in_sphere(search, radius, pc); }
            },
            Err(_) => {}
        }

        let (current_search, current_val) = match self.dimension {
            0 => (search.x(), self.val.x()),
            1 => (search.y(), self.val.y()),
            _ => (search.z(), self.val.z())
        };

        let border_left = current_search - radius;
        let border_right = current_search + radius;



        match comp {
            Ok(res) => match res {
                Ordering::Less => if let Some(ref node) = (&self).right {
                    if border_right >= current_val {
                        node.in_sphere(search, radius, pc);
                    }
                },
                Ordering::Greater => if let Some(ref node) = (&self).left {
                    if border_left <= current_val {
                        node.in_sphere(search, radius, pc);
                    }
                },
                Ordering::Equal => {}
            },
            Err(_) => {}
        }
    }

    pub fn in_box<PSearch>(&self, search: &PSearch, x_size: f64, y_size: f64, z_size: f64, pc: &mut Vec<PFind>) where
        PSearch: Is3D {

        if x_size <= 0.0 || y_size <= 0.0 || z_size <= 0.0 { return; }

        if let (Ok(dist_x), Ok(dist_y), Ok(dist_z)) = (dimension_dist(search, &self.val, 0), dimension_dist(search, &self.val, 1), dimension_dist(search, &self.val, 2)) {
            if dist_x <= 0.5 * x_size && dist_y <= 0.5 * y_size && dist_z <= 0.5 * z_size {
                pc.push(self.val.clone());
            }

            if self.is_leaf()  { return; }

            let comp = dimension_compare(search, &self.val, self.dimension);

            match comp {
                Ok(res) => match res {
                    Ordering::Less  => if let Some(ref node) = (&self).left { node.in_box(search, x_size, y_size, z_size, pc); },
                    _               => if let Some(ref node) = (&self).right { node.in_box(search, x_size, y_size, z_size, pc); }
                },
                Err(_) => {}
            }

            let (current_search, current_val, current_size) = match self.dimension {
                0 => (search.x(), self.val.x(), x_size),
                1 => (search.y(), self.val.y(), y_size),
                _ => (search.z(), self.val.z(), z_size)
            };

            let border_left = current_search - 0.5 * current_size;
            let border_right = current_search + 0.5 * current_size;

            match comp {
                Ok(res) => match res {
                    Ordering::Less => if let Some(ref node) = (&self).right {
                        if border_right >= current_val {
                            node.in_box(search, x_size, y_size, z_size, pc);
                        }
                    },
                    Ordering::Greater => if let Some(ref node) = (&self).left {
                        if border_left <= current_val {
                            node.in_box(search, x_size, y_size, z_size, pc);
                        }
                    },
                    Ordering::Equal => {}
                },
                Err(_) => {}
            }
        }
    }
}