nb-tree 0.2.0-alpha01

use itertools::Itertools;

use crate::{
    prelude::{Path, Tree},
    tree::NodeIDX,
};
use std::{any::Any, hash::Hash, iter::FusedIterator, marker::PhantomData, mem};

use super::depth::{Traversal, TreeIterTarget};

pub struct IterSortable<'a, N, B, T, S>
where
    T: TreeIterTarget<'a, N, B>,
    S: FnMut(&mut Vec<(&'a B, usize)>),
{
    tree: &'a Tree<N, B>,
    root: Option<NodeIDX>,
    stages: Vec<Vec<(&'a B, usize)>>,
    up: usize,
    sort: S,
    _phantom: PhantomData<T>,
}

impl<'a, N, B, T, S> IterSortable<'a, N, B, T, S>
where
    T: TreeIterTarget<'a, N, B>,
    S: FnMut(&mut Vec<(&'a B, usize)>),
{
    pub fn new(tree: &'a Tree<N, B>, sort: S) -> Self {
        Self {
            tree,
            root: tree.get_root_idx(),
            stages: vec![],
            up: 0,
            sort,
            _phantom: PhantomData,
        }
    }
}

impl<'a, N, B, T> IterSortable<'a, N, B, T, fn(&mut Vec<(&'a B, usize)>)>
where
    T: TreeIterTarget<'a, N, B>,
{
    pub fn new_unsorted(tree: &'a Tree<N, B>) -> Self {
        Self {
            tree,
            root: tree.get_root_idx(),
            stages: vec![],
            up: 0,
            sort: |_| {},
            _phantom: PhantomData,
        }
    }

    pub fn up(&self) -> usize {
        self.up
    }
}

impl<'a, N, B, T, S> IterSortable<'a, N, B, T, S>
where
    B: Clone + Eq + Hash,
    T: TreeIterTarget<'a, N, B>,
    S: FnMut(&mut Vec<(&'a B, usize)>),
{
    pub fn new_sub_state(
        tree: &'a Tree<N, B>,
        path: Path<B>,
        sort: S,
    ) -> Result<Self, Option<Path<B>>> {
        Ok(Self {
            tree,
            root: Some(
                tree.get_idx(&path, None)
                    .map_err(|r| r.map(|(_, idx_p)| path.path_to(idx_p)))?,
            ),
            stages: vec![],
            up: 0,
            sort,
            _phantom: PhantomData,
        })
    }
}

impl<'a, N, B, T> IterSortable<'a, N, B, T, fn(&mut Vec<(&'a B, usize)>)>
where
    B: Clone + Eq + Hash,
    T: TreeIterTarget<'a, N, B>,
{
    pub fn new_sub_state_unsorted(
        tree: &'a Tree<N, B>,
        path: Path<B>,
    ) -> Result<Self, Option<Path<B>>> {
        Ok(Self {
            tree,
            root: Some(
                tree.get_idx(&path, None)
                    .map_err(|r| r.map(|(_, idx_p)| path.path_to(idx_p)))?,
            ),
            stages: vec![],
            up: 0,
            sort: |_| {},
            _phantom: PhantomData,
        })
    }

    pub fn is_last_child(&mut self) -> bool {
        if let Some(last) = self.stages.last_mut() {
            last.is_empty()
        } else {
            true
        }
    }

    /// Tries to visit the given branch next.
    pub fn try_follow(&mut self, branch: B) -> bool {
        if let Some(children) = self.stages.last_mut() {
            if let Some((idx, _)) = children.iter().find_position(|&&x| *x.0 == branch) {
                let last = children.len() - 1;
                children.swap(idx, last);
                true
            } else {
                false
            }
        } else {
            false
        }
    }
}

impl<'a, N, B, T, S> Iterator for IterSortable<'a, N, B, T, S>
where
    T: TreeIterTarget<'a, N, B>,
    S: FnMut(&mut Vec<(&'a B, usize)>),
{
    type Item = Traversal<T::Target, &'a B>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(last) = self.stages.last_mut() {
            // Iterate on a child (not the root)
            if let Some((branch, node_idx)) = last.pop() {
                // Go down to a child
                let node = &self.tree.nodes[node_idx];
                self.stages
                    .push(node.children.iter().map(|(b, i)| (b, *i)).collect());
                //self.path.push_leaf(branch.clone());
                Some(Traversal::Step {
                    up: std::mem::replace(&mut self.up, 0),
                    branch,
                    data: T::target(self.tree, node_idx),
                })
                //Some((self.path.clone(), &node.value))
            } else {
                // Move back up to the parent
                self.stages.pop();
                //self.path.pop_leaf();
                self.up += 1;
                self.next()
            }
        } else if let Some(root) = self.root.take() {
            let node = &self.tree.nodes[root];
            self.stages
                .push(node.children.iter().map(|(b, i)| (b, *i)).collect());
            Some(Traversal::Start(T::target(self.tree, root)))
        } else {
            None
        }
    }
}

impl<'a, N, B, T, S> FusedIterator for IterSortable<'a, N, B, T, S>
where
    T: TreeIterTarget<'a, N, B>,
    S: FnMut(&mut Vec<(&'a B, usize)>),
{
}