use ptree::TreeBuilder;
use std::cmp::Ordering;

// Min. order: 3
const ORDER: usize = 4;

#[derive(Debug)]
struct Node<T> {
    entries: [Option<T>; ORDER - 1],
    children: [Option<usize>; ORDER],
    parent: Option<usize>,
}

impl<T: Ord> Node<T> {
    fn new(parent: Option<usize>) -> Node<T> {
        Node {
            entries: Default::default(),
            children: Default::default(),
            parent: parent,
        }
    }
}

#[derive(Debug)]
struct BTree<T> {
    v: Vec<Node<T>>,
    root: usize,
}

impl<T: std::fmt::Debug + Ord> BTree<T> {
    fn new(initial_capacity: usize) -> BTree<T> {
        BTree {
            v: Vec::with_capacity(initial_capacity),
            root: 0,
        }
    }

    fn free(&self, node: usize) -> usize {
        self.v[node]
            .entries
            .iter()
            .map(|e| if e.is_some() { 0 } else { 1 })
            .sum()
    }
    fn get(&self, node: usize, elem: usize) -> &Option<T> {
        &self.v[node].entries[elem]
    }
    fn get_mut(&mut self, node: usize, elem: usize) -> &mut Option<T> {
        &mut self.v[node].entries[elem]
    }
    fn get_child(&self, node: usize, child: usize) -> &Option<usize> {
        &self.v[node].children[child]
    }
    fn get_child_mut(&mut self, node: usize, child: usize) -> &mut Option<usize> {
        &mut self.v[node].children[child]
    }

    /// Insert a value into an existing node, preserving the order.
    /// Requires that at least one slot be free.
    fn insert_inline(
        &mut self,
        node: usize,
        value: T,
        left: Option<usize>,
        right: Option<usize>,
    ) -> usize {
        assert!(self.free(node) >= 1);
        let nodeix = node;

        let mut loc = ORDER - 2;

        let node = &mut self.v[node];
        for i in 0..ORDER - 1 {
            if let Some(ref el) = node.entries[i] {
                if value.cmp(el) == Ordering::Less {
                    // Found right location. Shift remaining elements.
                    loc = i;
                    break;
                }
            } else {
                loc = i;
                break;
            }
        }

        for i in 0..(ORDER - 2 - loc) {
            node.entries[ORDER - 2 - i] = node.entries[ORDER - 3 - i].take();
            node.children[ORDER - 1 - i] = node.children[ORDER - 2 - i].take();
        }
        node.entries[loc] = Some(value);
        node.children[loc] = left;
        node.children[loc + 1] = right;

        if let Some(left) = left {
            self.v[left].parent = Some(nodeix)
        }
        if let Some(right) = right {
            self.v[right].parent = Some(nodeix);
        }
        loc
    }

    /// Split a node, returning the IDs of the two new nodes and the pivot value.
    /// Parent pointers need to be updated by caller.
    fn split(&mut self, node: usize) -> (usize, T, usize) {
        assert!(self.free(node) <= 1);
        // `node` will become left node.
        let leftix = node;
        let rightix = self.v.len();

        println!("Starting split: {:?}", self.v[node]);

        let mut right = Node::<T>::new(None); // Parent updated by caller!

        let pivotix = (ORDER - 1) / 2;
        let pivot = self.get_mut(node, pivotix).take();

        // Pivot-left child remains in left node.
        for i in pivotix + 1..ORDER {
            right.children[i - pivotix - 1] = self.get_child_mut(leftix, i).take();
        }
        // Make sure to move all children pointers!
        //right.children[ORDER - pivotix] = self.get_child_mut(leftix, ORDER - 1).take();
        for i in pivotix + 1..ORDER - 1 {
            right.entries[i - pivotix - 1] = self.get_mut(leftix, i).take();
        }

        self.v.push(right);
        println!(
            "Finishing split: left {:?} right {:?} pivot {:?}",
            self.v[leftix], self.v[rightix], pivot
        );
        (leftix, pivot.unwrap(), rightix)
    }

    fn is_leaf(&self, node: usize) -> bool {
        self.v[node].children.iter().all(|c| c.is_none())
    }
    fn update_right_child(&mut self, node: usize, ch: usize) {
        for i in 0..ORDER - 1 {
            if self.get(node, i).is_none() {
                if let Some(subch) = self.get_child(node, i) {
                    self.update_right_child(*subch, ch);
                    return;
                } else {
                    *self.get_child_mut(node, i) = Some(ch);
                    return;
                }
            }
        }
        // This means that even the last child slot was used so far.
        self.update_right_child(self.get_child(node, ORDER - 1).unwrap(), ch);
    }

    /// Insert & root splitting.
    fn insert(&mut self, value: T) -> bool {
        if let Some((l, pivot, r)) = self.insert_into(self.root, value) {
            // Split root
            println!("Splitting root with new children {}, {}", l, r);
            let mut newroot = Node::new(None);
            let rootix = self.v.len();
            newroot.children[0] = Some(l);
            newroot.children[1] = Some(r);
            newroot.entries[0] = Some(pivot);

            self.v[l].parent = Some(rootix);
            self.v[r].parent = Some(rootix);
            self.root = rootix;

            self.v.push(newroot);
            true
        } else {
            true
        }
    }

    /// Insert value recursively. If a split occurs, returns (left, pivot, right) tuple.
    fn insert_into(&mut self, node: usize, value: T) -> Option<(usize, T, usize)> {
        // BTree initialized?
        if self.v.len() == 0 {
            println!("Initializing");
            self.v.push(Node::new(None));
            self.insert_inline(self.root, value, None, None);
            return None;
        }
        // Leaf node with free space? We are finished.
        if self.is_leaf(node) && self.free(node) >= 1 {
            println!("Leaf {} with free space", node);
            self.insert_inline(node, value, None, None);
            return None;
        } else if self.is_leaf(node) && self.free(node) < 1 {
            // Leaf node but needs to be split.
            let (l, pivot, r) = self.split(node);
            // Full leaf node?
            println!("Splitting leaf {} into {} {:?} {}", node, l, pivot, r);

            // Insert value into left or right subtree.
            match value.cmp(&pivot) {
                Ordering::Less => assert!(self.insert_into(l, value).is_none()),
                Ordering::Greater => assert!(self.insert_into(r, value).is_none()),
                _ => panic!("Attempting to insert duplicate element"),
            };
            return Some((l, pivot, r));
        } else {
            // Descend into node.
            println!("Scanning node {}", node);

            let mut split = None;
            // Find correct child.
            for i in 0..ORDER - 1 {
                if i < ORDER - 2 {
                    if let Some(el) = self.get(node, i) {
                        if value.cmp(el) == Ordering::Less {
                            split = self.insert_into(self.get_child(node, i).unwrap(), value);
                            break;
                        }
                    } else {
                        if let Some(ch) = self.get_child(node, i) {
                            split = self.insert_into(*ch, value);
                            break;
                        }
                        panic!(
                            "Internal node missing child? {} => {:?}",
                            node, self.v[node]
                        );
                    }
                } else {
                    if let Some(ch) = self.get_child(node, i) {
                        split = self.insert_into(*ch, value);
                        break;
                    } else {
                        let (l, pivot, r) = self.split(node);

                        // Insert value into left or right subtree.
                        match value.cmp(&pivot) {
                            Ordering::Less => assert!(self.insert_into(l, value).is_none()),
                            Ordering::Greater => assert!(self.insert_into(r, value).is_none()),
                            _ => panic!("Attempting to insert duplicate element"),
                        };

                        return Some((l, pivot, r));
                    }
                }
            }

            println!("Split necessary? {:?}", split);
            // Merge split if one happened below us.
            if let Some((l, pivot, r)) = split {
                if self.free(node) >= 1 {
                    // We can insert the split here.
                    self.insert_inline(node, pivot, Some(l), Some(r));
                    return None;
                } else {
                    // Split this node too.
                    let (nl, npivot, nr) = self.split(node);
                    match pivot.cmp(&npivot) {
                        Ordering::Less => self.insert_inline(nl, pivot, Some(l), Some(r)),
                        Ordering::Greater => self.insert_inline(nr, pivot, Some(l), Some(r)),
                        _ => panic!("Attempting to insert duplicate element"),
                    };
                    return Some((nl, npivot, nr));
                }
            } else {
                return None;
            }
        }
    }
}
impl<T: std::fmt::Debug> BTree<T> {
    fn debug_print(&self) -> String {
        let mut t = TreeBuilder::new(format!("<{}>", self.root));
        self.print_subtree(&mut t, self.root);
        let mut o: Vec<u8> = Vec::new();
        ptree::write_tree(&t.build(), &mut o);
        String::from_utf8(o).unwrap()
    }
    fn print_subtree(&self, tb: &mut TreeBuilder, node: usize) {
        let node = &self.v[node];
        for i in 0..ORDER {
            if let Some(ch) = node.children[i] {
                self.print_subtree(tb.begin_child(format!("<{}>", ch)), ch);
                tb.end_child();
            }
            if i < ORDER - 1 {
                tb.add_empty_child(format!("{:?}", node.entries[i]));
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::BTree;
    #[test]
    fn basic_test() {
        let mut bt = BTree::<i32>::new(16);

        let ns = &[15, 2, 21, 35, 4, 62, 31, 27, 88, 1, 12, 45, 11, 16, 23];

        for n in ns {
            bt.insert(*n);
            println!("Inserting {}", n);
            println!("root: {}", bt.root);
            println!("{}", bt.debug_print());
        }
        for (i, e) in bt.v.iter().enumerate() {
            println!("{} {:?}", i, e);
        }
    }
}