gst 0.1.3

use arrayvec::ArrayVec;
use std::convert::Into;

/// `GstKey` represents a the Generalized Search Tree.
/// See http://gist.cs.berkeley.edu/libgist-2.0/
/// The rename from `GiST` to Gst is because it conflicts with github gists and no one knows what the
/// 'i' stands for anyway.
pub trait GstKey: Sized + Clone + Eq {
    /// Return true if self contains (or overlaps) k.
    fn consistent(&self, k: &Self) -> bool;

    /// Expand the Key to cover k.
    fn expand(&self, k: &Self) -> Self;

    /// Given a slice of T, make a Self that contains them. Return None if the slice was empty.
    fn union(items: &[Self]) -> Option<Self> {
        let mut iter = items.iter();
        let first = iter.next();
        match first {
            None => None,
            Some(init) => Some(iter.fold(init.clone().into(), |lhs, rhs| lhs.expand(rhs)))
        }
    }

    // Compress the Key. The Compression can be lossy.
    //fn compress(&mut self);

    // Decompress the Key. The compression can be lossy.
    //fn decompress(&mut self);

    /// Return a float representing the cost of inserting T into the Key.
    /// Also return a score so the machinations can compare scores from orphaned children.
    fn penalty(bounds: &[Self], t: &Self) -> (usize, f32);

    /// Given a slice of Self, return a vector of indices of items that should be carved off into
    /// their own node.
    /// Assumes that the leaf was inserted in some kind of ordering.
    fn pick_split(bounds: &[Self], min_split_size: usize) -> Vec<usize>;
}

const B: usize = 6;
const CAPACITY: usize = B*2-1;
const MIN_OCCUPANCY: usize = CAPACITY / 3;

/// The Error type
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum Error{
    NotFound
}

/// `Gst` is a tree similar to a B+ tree but using a `GstKey` instead of requiring that the key
/// implements `Ord`
pub struct Gst<K: GstKey, V> {
    root: Box<Node<K, V>>,
}

impl<K: GstKey, V> Gst<K, V> {
    /// Create a new Gst
    pub fn new() -> Self {
        Gst {
            root: Box::new(Node::Internal(InternalNode::new())),
        }
    }

    /// Insert a Key and associated Value into the Gst.
    pub fn insert(&mut self, key: K, value: V) {
        let spill = self.root.insert(NodeType::Root, key, value);
        assert!(spill.is_none());
    }

    /// Retrieve all the keys and values inside a region.
    pub fn get(&self, key: &K) -> Vec<(&K, &V)> {
        self.root.get(&key)
    }

    /// Remove an item based on an exact key match.
    /// Returns the Key and Value on success, otherwise `Error::NotFound`
    pub fn remove(&mut self, key: &K) -> Result<(K, V), Error> {
        self.root.remove(&key)
    }
}

enum NodeType { Root, Internal }

impl NodeType {
    fn is_root(&self) -> bool {
        match *self {
            NodeType::Root => true,
            _ => false
        }
    }
}

enum Node<K: GstKey, V> {
    Leaf(LeafNode<K, V>),
    Internal(InternalNode<K, V>)
}

impl<K: GstKey, V> Node<K, V> {
    fn bounds(&self) -> K {
        match *self {
            Node::Leaf(ref leaf) => K::union(&leaf.keys).expect("Unpruned leaf has no bounding box."),
            Node::Internal(ref internal) => K::union(&internal.bounds).unwrap()
        }
    }

    fn get(&self, key: &K) -> Vec<(&K, &V)>{
        match *self {
            Node::Leaf(ref leaf) => leaf.get(&key),
            Node::Internal(ref internal) => internal.get(&key)
        }
    }

    fn insert(&mut self, nodetype: NodeType, key: K, value: V) -> Option<Box<Node<K, V>>> {
        match *self {
            Node::Leaf(ref mut leaf) => leaf.insert(key, value),
            Node::Internal(ref mut internal) => internal.insert(nodetype, key, value)
        }
    }

    fn remove(&mut self, key: &K) -> Result<(K, V), Error> {
        match *self {
            Node::Leaf(ref mut leaf) => leaf.remove(key),
            Node::Internal(ref mut internal) => internal.remove(key)
        }
    }
}

impl<K: GstKey, V> From<InternalNode<K, V>> for Node<K, V> {
    fn from(node: InternalNode<K, V>) -> Self {
        Node::Internal(node)
    }
}

impl<K: GstKey, V> From<LeafNode<K, V>> for Node<K, V> {
    fn from(node: LeafNode<K, V>) -> Self {
        Node::Leaf(node)
    }
}

struct LeafNode<K: GstKey, V> {
    keys: ArrayVec<[K; CAPACITY]>,
    values: ArrayVec<[V; CAPACITY]>
}

impl<K: GstKey, V> LeafNode<K, V> {
    fn new() -> Self {
        LeafNode{ keys: ArrayVec::new(), values: ArrayVec::new() }
    }

    fn new_with(key: K, value: V) -> Self {
        let mut n = LeafNode::new();
        n.keys.push(key);
        n.values.push(value);
        n
    }

    fn get(&self, key: &K) -> Vec<(&K, &V)> {
        let mut res = vec![];
        for i in 0..self.keys.len() {
            if key.consistent(&self.keys[i]) {
                res.push((&self.keys[i], &self.values[i]));
            }
        }
        res
    }

    fn insert(&mut self, key: K, value: V) -> Option<Box<Node<K, V>>> {
        if self.keys.len() < CAPACITY {
            self.keys.push(key);
            self.values.push(value);
            None
        } else {
            let mut orphan = self.split();
            // wrong. need to compare between self and orphan.
            orphan.keys.push(key);
            orphan.values.push(value);
            Some(Box::new(orphan.into()))
        }
    }

    fn split(&mut self) -> LeafNode<K, V> {
        let split_indices = K::pick_split(&self.keys, MIN_OCCUPANCY);
        let mut orphan = Self::new();
        for (offset, idx) in split_indices.into_iter().enumerate() {
            let real_idx = idx - offset;
            let key = self.keys.drain(real_idx..real_idx+1).next().unwrap();
            orphan.keys.push(key);
            let value = self.values.drain(real_idx..real_idx+1).next().unwrap();
            orphan.values.push(value);
        }
        orphan
    }

    fn remove(&mut self, key: &K) -> Result<(K, V), Error> {
        let mut id = None;
        for (i, k) in self.keys.iter().enumerate() {
            if *key == *k {
                id = Some(i)
            }
        }

        if let Some(id) = id {
            let drained_key = self.keys.drain(id..id+1).next().unwrap();
            let drained_val= self.values.drain(id..id+1).next().unwrap();
            Ok((drained_key, drained_val))
        } else {
            Err(Error::NotFound)
        }
    }
}

struct InternalNode<K: GstKey, V> {
    bounds: ArrayVec<[K; CAPACITY]>,
    children: ArrayVec<[Box<Node<K, V>>; CAPACITY]>,
}

impl<K: GstKey, V> InternalNode<K, V> {
    fn new() -> Self {
        InternalNode{ bounds: ArrayVec::new(), children: ArrayVec::new() }
    }

    fn get(&self, key: &K) -> Vec<(&K, &V)> {
        let mut res = vec![];
        for (child, bounds) in self.children.iter().zip(self.bounds.iter()) {
            if bounds.consistent(&key) {
                res.append(&mut child.get(&key));
            }
        }
        res
    }


    fn insert(&mut self, nodetype: NodeType, key: K, value: V) -> Option<Box<Node<K, V>>> {
        if self.children.len() == 0 {
            let leaf = Box::new(LeafNode::new_with(key, value).into());
            self.insert_node(leaf);
            None
        } else if self.children.len() == CAPACITY {
            if nodetype.is_root() {
                self.split_root();
                let spill = self.insert(nodetype, key, value); // try again
                assert!(spill.is_none());
                spill
            } else {
                let mut orphan = self.split();
                let (orphan_idx, orphan_score) = K::penalty(&orphan.bounds, &key);
                // fast insert: if penalty is 0, use it.
                if orphan_score == 0. {
                    orphan.insert_node_and_spill(orphan_idx, key, value);
                } else {
                    let (self_idx, self_score) = K::penalty(&self.bounds, &key);
                    if orphan_score < self_score {
                        orphan.insert_node_and_spill(orphan_idx, key, value);
                    } else {
                        self.insert_node_and_spill(self_idx, key, value);
                    }
                }
                Some(Box::new(orphan.into()))
            }
        } else {
            let (idx, _) = K::penalty(&self.bounds, &key);
            self.insert_node_and_spill(idx, key, value);
            None
        }
    }

    fn insert_node_and_spill(&mut self, idx: usize, key: K, value: V) {
        let spill = self.children[idx].insert(NodeType::Internal, key, value);
        self.bounds[idx] = self.children[idx].bounds();
        if let Some(node) = spill {
            self.insert_node(node.into());
        }
    }

    fn insert_node(&mut self, node: Box<Node<K, V>>) {
        assert!(self.children.len() < CAPACITY);
        let bounds = node.bounds();
        self.children.push(node);
        self.bounds.push(bounds);
    }

    fn split(&mut self) -> Self {
        let split_indices = K::pick_split(&self.bounds, MIN_OCCUPANCY);
        let mut orphan = Self::new();
        for (offset, idx) in split_indices.into_iter().enumerate() {
            let real_idx = idx - offset;
            let child = self.children.drain(real_idx..real_idx + 1).next().unwrap();
            orphan.insert_node(child);

            // remove from bounds as well
            self.bounds.drain(real_idx..real_idx+1).next();
        }
        orphan
    }

    fn split_root(&mut self) {
        let split_indices = K::pick_split(&self.bounds, MIN_OCCUPANCY);
        let mut lhs = Self::new();
        for (offset, idx) in split_indices.iter().enumerate() {
            let real_idx = *idx - offset;
            let child = self.children.drain(real_idx..real_idx + 1).next().unwrap();
            lhs.insert_node(child);
        }

        let mut rhs = Self::new();
        for child in self.children.drain(..) {
            rhs.insert_node(child);
        }

        self.bounds.clear();

        assert!(self.children.len() == 0, "Root node not empty after split.");
        self.insert_node(Box::new(lhs.into()));
        self.insert_node(Box::new(rhs.into()));
    }

    fn remove(&mut self, key: &K) -> Result<(K, V), Error> {
        for (child, bounds) in self.children.iter_mut().zip(self.bounds.iter()) {
            if bounds.consistent(&key) {
                let res = child.remove(&key);
                if res.is_ok() {
                    return res;
                }
            }
        }
        Err(Error::NotFound)
    }

}

pub mod diag {

use super::{Gst, Node};
use rtree::Rect;
fn print_rect(bound: &Rect, color: &str) {
    println!("<rect x=\"{}\" y=\"{}\" width=\"{}\" height=\"{}\" style=\"stroke:{};stroke-width:0.0001;fill-opacity:0.1;stroke-opacity:0.9\"/>"
             , bound.xmin.0, bound.ymin.0, bound.xmax.0-bound.xmin.0, bound.ymax.0-bound.ymin.0, color);
}

pub fn print_gst<V>(tree: &Gst<Rect, V>) {
    println!("<svg xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\" width=\"800\" height=\"600\" viewBox=\"5 15 47 56\">");
    print_node(&tree.root);
    println!("</svg>");
}

fn print_node<V>(node: &Box<Node<Rect, V>>) {
    match **node {
        Node::Leaf(ref leaf) => {
            for key in leaf.keys.iter() {
                print_rect(key, "blue");
            }
        },
        Node::Internal(ref internal) => {
            for bound in internal.bounds.iter() {
                print_rect(bound, "red");
            }
            for child in internal.children.iter() {
                print_node(child);
            }
        }
    }
}

}

#[cfg(test)]
mod test {
    use gst::{LeafNode, InternalNode, NodeType, Error};
    use rtree::{Point, Rect};

    #[test]
    fn test_leaf_get() {
        let leaf = LeafNode::new_with(Rect::from_float(0., 0., 0., 0.), String::from("Origin"));
        assert_eq!(leaf.get(&Rect::from_float(0.0,  0.5,  0.0,  0.5 )).len(), 1);
    }

    #[test]
    fn test_leaf_split() {
        let mut leaf = LeafNode::new();
        for i in 0..10 {
            let key = Rect::from(Point::new(i as f32, i as f32).into());
            let val = format!("{}", i*i);
            assert!(leaf.insert(key, val).is_none());
        }
        let spill = leaf.split();
        assert_eq!(leaf.keys.len(), 5);
        assert_eq!(spill.keys.len(), 5);
        assert_eq!(leaf.get(&Rect::from_float(0.0, 10.0,  0.0, 10.0 )).len(), 5);
        assert_eq!(spill.get(&Rect::from_float(0.0, 10.0,  0.0, 10.0 )).len(), 5);
    }

    #[test]
    fn test_leaf_remove() {
        let rect = Rect::from_float(0., 0., 0., 0.);
        let string = String::from("Origin");
        let mut leaf = LeafNode::new_with(rect, string.clone());
        assert_eq!(leaf.remove(&Rect::from_float(0.0,  0.5,  0.0,  0.5 )), Err(Error::NotFound));
        assert_eq!(leaf.get(&Rect::from_float(0.0,  0.5,  0.0,  0.5 )).len(), 1);
        assert_eq!(leaf.remove(&rect), Ok((rect, string.clone())));
        assert_eq!(leaf.get(&Rect::from_float(0.0,  0.5,  0.0,  0.5 )).len(), 0);
    }

    #[test]
    fn test_internal_get() {
        let mut internal = InternalNode::new();
        for i in 0..10 {
            let point = Point::new(i as f32, i as f32);
            let leaf = LeafNode::new_with(Rect::from(point.into()), format!("{}", i*i));
            internal.insert_node(Box::new(leaf.into()));
        }
        assert_eq!(internal.get(&Rect::from_float(0.0,  0.5,  0.0,  0.5 )).len(), 1);
        assert_eq!(internal.get(&Rect::from_float(0.0,  1.0,  0.0,  1.0 )).len(), 2);
        assert_eq!(internal.get(&Rect::from_float(0.0, 10.0,  0.0, 10.0 )).len(), 10);
    }

    #[test]
    fn test_internal_remove() {
        let mut internal = InternalNode::new();
        for i in 0..10 {
            let point = Point::new(i as f32, i as f32);
            let leaf = LeafNode::new_with(Rect::from(point.into()), format!("{}", i*i));
            internal.insert_node(Box::new(leaf.into()));
        }
        assert_eq!(internal.remove(&Rect::from_float(0.0,  0.5,  0.0,  0.5 )), Err(Error::NotFound));
        assert!(internal.remove(&Rect::from_float(1.,  1.,  1.,  1.)).is_ok());
    }


    #[test]
    fn test_internal_root_get_split() {
        let mut internal = InternalNode::new();
        for i in 0..20 {
            let key = Point::new(i as f32, i as f32);
            let val = format!("{}", i*i);
            assert!(internal.insert(NodeType::Root, key.into(), val).is_none());
        }
        assert_eq!(internal.get(&Rect::from_float(0.0,  0.5,  0.0,  0.5 )).len(), 1);
        assert_eq!(internal.get(&Rect::from_float(0.0,  1.0,  0.0,  1.0 )).len(), 2);
        assert_eq!(internal.get(&Rect::from_float(0.0, 20.0,  0.0, 20.0 )).len(), 20);
    }

    #[test]
    fn test_internal_split_root() {
        let mut internal = InternalNode::new();
        for i in 0..10 {
            let point = Point::new(i as f32, i as f32);
            let leaf = LeafNode::new_with(Rect::from(point.into()), format!("{}", i*i));
            internal.insert_node(Box::new(leaf.into()));
        }
        assert_eq!(internal.children.len(), 10);
        internal.split_root();
        assert_eq!(internal.children.len(), 2);
    }

    #[test]
    fn test_internal_split() {
        let mut internal = InternalNode::new();
        for i in 0..10 {
            let point = Point::new(i as f32, i as f32);
            let leaf = LeafNode::new_with(Rect::from(point.into()), format!("{}", i*i));
            internal.insert_node(Box::new(leaf.into()));
        }
        assert_eq!(internal.children.len(), 10);
        let spill = internal.split();
        println!("internal {:?}", internal.children.len());
        println!("spill {:?}", spill.children.len());
        assert_eq!(internal.children.len() + spill.children.len(), 10);
        assert!(internal.children.len() > 0);
        assert!(spill.children.len() > 0);

        assert_eq!(internal.get(&Rect::from_float(0.0,  10.0,  0.0,  10.0 )).len(), internal.children.len());
        assert_eq!(spill.get(&Rect::from_float(0.0,  10.0,  0.0,  10.0 )).len(), spill.children.len());
    }
}