use indextree::{Arena, NodeId};

trait CalcNode {
    fn update(&mut self, input: Vec<&dyn CalcNode>);
    fn output(&self) -> i64;
}

struct NodeModel {
    output: i64,
    internal: i64,
}

impl CalcNode for NodeModel {
    //input can't be cloned/copied in final version
    fn update(&mut self, input: Vec<&dyn CalcNode>) {
        //example for n possible update funciton
        let chid_input: i64 = input.iter().map(|&f| f.output()).sum();
        self.output = self.internal * chid_input;
        self.internal += 1
    }

    fn output(&self) -> i64 {
        // output is not copy-able in final version
        self.output
    }
}

impl NodeModel {
    fn new(internal: i64) -> Self {
        Self {
            output: internal,
            internal,
        }
    }
}

fn main() {
    let (mut arena_, root) = build_tree();
    let arena = &mut arena_;

    //update 3 times as test
    let node_stack: Vec<NodeId> = root.descendants(arena).collect(); //collect all nodes depth first

    for _ in 0..3 {
        //update in reverse order deepest node first, ROOT LAST.
        for &n in node_stack.iter().rev() {
            // collect updated children when available
            let children: Vec<&dyn CalcNode> = {
                n.children(arena)
                    .map(|f| arena.get(f).unwrap().get().as_ref())
                    .collect()
            };

            //get the node to update
            let node = { arena.get_mut(n).unwrap().get_mut().as_mut() }; //<- can't borrow arena here to get the node

            // update the node
            node.update(children);
        }
    }
}

fn build_tree() -> (Arena<Box<dyn CalcNode>>, NodeId) {
    // Create a new arena
    let mut arena: Arena<Box<dyn CalcNode>> = Arena::new();

    // Add some new nodes to the arena
    let a = arena.new_node(Box::new(NodeModel::new(1)));
    let b = arena.new_node(Box::new(NodeModel::new(2)));
    let c = arena.new_node(Box::new(NodeModel::new(3)));
    let d = arena.new_node(Box::new(NodeModel::new(4)));
    let e = arena.new_node(Box::new(NodeModel::new(5)));
    let f = arena.new_node(Box::new(NodeModel::new(6)));
    let g = arena.new_node(Box::new(NodeModel::new(7)));
    let h = arena.new_node(Box::new(NodeModel::new(8)));

    // Build tree
    //           a
    //        b     c
    //       d e    f
    //             g h
    a.append(b, &mut arena);
    a.append(c, &mut arena);
    b.append(d, &mut arena);
    b.append(e, &mut arena);
    c.append(f, &mut arena);
    f.append(g, &mut arena);
    f.append(h, &mut arena);

    return (arena, a);
}