extern crate num;

use std::fmt;
use std::ops::Index;
use std::cmp::PartialOrd;

pub trait Zero {
    fn zero() -> Self;
}

impl<T: num::Zero> Zero for T {
    fn zero() -> Self {
        num::Zero::zero()
    }
}

/// A k-dimension kd tree
pub enum KdTree<T> {
    Node { location: T, left:  Box<KdTree<T>>, right: Box<KdTree<T>> },
    Leaf(Vec<T>),
}

impl<T, U> KdTree<T> where T: Copy + Clone + Index<usize, Output = U> + Zero, U: PartialOrd {
    pub fn new(points: &mut [T], dimension: usize, leaf_size: usize) -> KdTree<T> {
        KdTree::new_with_depth(points, dimension, leaf_size, 0)
    }

    fn new_with_depth(points: &mut [T], dimension: usize, leaf_size: usize, depth: usize) -> KdTree<T> {

        // Select axis to work on
        let axis = depth % dimension;

        // Sort the points by axis
        points.sort_by(|p1, p2| p1[axis].partial_cmp(&p2[axis]).unwrap());

        // Find the median
        let len = points.len();

        // If there are not enough elements for a leaf, make a leaf
        if len <= leaf_size {
            let mut contained = vec![Zero::zero(); len];
            contained.copy_from_slice(points);
            return KdTree::Leaf(contained);
        }

        // Else split what's remaining
        let delimiter = points.len() / 2;
        let median = points[delimiter];

        // Build the left and right branches
        let zero = Zero::zero();
        let mut left = vec![zero; delimiter];
        left.copy_from_slice(&points[0..delimiter]);

        let right = &mut points[delimiter + 1 .. len];

        // Recursive call
        KdTree::Node {
            location: median,
            left:  Box::new(KdTree::new_with_depth(&mut left, dimension, leaf_size, depth + 1)),
            right: Box::new(KdTree::new_with_depth(right, dimension, leaf_size, depth + 1)),
        }
    }
}

impl<T: Zero + Copy + Clone + fmt::Display + Index<usize>> KdTree<T> {
    fn print(&self, formatter: &mut fmt::Formatter, indent: &str) -> fmt::Result {
        match self {
            &KdTree::Node { location, ref left, ref right } => {
                writeln!(formatter, "{}{}", indent, location)?;
                left.print(formatter, &format!("{}  ", indent))?;
                right.print(formatter, &format!("{}  ", indent))?;
            },

            &KdTree::Leaf(ref elements) => {
                write!(formatter, "{}[", indent)?;
                for element in elements {
                    write!(formatter, "{}, ", element)?;
                }
                writeln!(formatter, "]")?;
            },
        }

        Ok(())
    }
}

impl<T: Zero + Copy + Clone + fmt::Display + Index<usize>> fmt::Display for KdTree<T> {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        self.print(formatter, "")
    }

}