// https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=7cccabd269fd1ee8f61ff23fd79117e7

#[derive(Debug)]
struct Node<T>
where
    T: PartialEq,
{
    idx: usize,
    val: T,
    parent: Option<usize>,
    children: Vec<usize>,
}

impl<T> Node<T>
where
    T: PartialEq,
{
    fn new(idx: usize, val: T) -> Self {
        Self {
            idx,
            val,
            parent: None,
            children: vec![],
        }
    }
}

#[derive(Debug, Default)]
struct ArenaTree<T>
where
    T: PartialEq,
{
    arena: Vec<Node<T>>,
}

impl<T> ArenaTree<T>
where
    T: PartialEq,
{
    fn node(&mut self, val: T) -> usize {
        //first see if it exists
        for node in &self.arena {
            if node.val == val {
                return node.idx;
            }
        }
        // Otherwise, add new node
        let idx = self.arena.len();
        self.arena.push(Node::new(idx, val));
        idx
    }
    fn size(&self) -> usize {
        self.arena.len()
    }
    fn edges(&self) -> usize {
        self.arena
            .iter()
            .fold(0, |acc, node| acc + node.children.len())
    }

    fn depth(&self, idx: usize) -> usize {
        match self.arena[idx].parent {
            Some(id) => 1 + self.depth(id),
            None => 0,
        }
    }
    fn depth_to_target(&self, idx: usize, target: &T) -> Option<usize> {
        // are we here?  If so, Some(0)
        if target == &self.arena[idx].val {
            return Some(0);
        }

        // If not, try all children
        for p in &self.arena[idx].children {
            if let Some(x) = self.depth_to_target(*p, &target) {
                return Some(1 + x);
            }
        }
        // If it cant be found, return None
        None
    }
    fn distance_between(&mut self, from: T, target: T) -> usize {
        // If it's not in the tree, this will add a new unconnected node
        // the final function will still return None
        let start_node = self.node(from);
        let mut ret = 0;
        // Start traversal
        let mut trav = &self.arena[start_node];
        // Explore all children, then hop up one
        while let Some(inner) = trav.parent {
            if let Some(x) = self.depth_to_target(inner, &target) {
                ret += x;
                break;
            }
            trav = &self.arena[inner];
            ret += 1;
        }
        // don't go all the way to target, just orbit
        if ret > 0 {
            ret - 1
        } else {
            ret
        }
    }
}

fn main() {
    let mut tree: ArenaTree<String> = ArenaTree::default();
    let hello = tree.node("Hello".into());
    let world = tree.node("World".into());
    tree.arena[hello].children.push(world);
    tree.arena[world].parent = Some(hello);
    let exclamation = tree.node("!".into());
    tree.arena[exclamation].parent = Some(world);
    tree.arena[world].children.push(exclamation);
    let period = tree.node(".".into());
    let period_2 = tree.node(".".into());
    assert_eq!(period, period_2);
    tree.arena[world].children.push(period);
    tree.arena[period_2].parent = Some(world);

    println!(
        "Total nodes: {}\nTotal edges: {}\nDepth of '.': {}\nDistance between World and !: {}",
        tree.size(),
        tree.edges(),
        tree.depth(period),
        tree.distance_between("World".into(), "!".into())
    );
}