nb-tree 0.2.0-alpha01

use std::collections::hash_map::Iter as HMIter;
use std::{hash::Hash, iter::FusedIterator, marker::PhantomData};

use crate::{path::Path, prelude::DiffNode, tree::node::TreeNode};

use super::super::{NodeIDX, Tree};

///////////////////////////////////////////////////////////////////////////////
//                                   Iter                                    //
///////////////////////////////////////////////////////////////////////////////

/// Describes a node's position relatively to the previous one and contains some data.
#[derive(Debug, PartialEq)]
pub enum Traversal<N, B> {
    Start(N),
    Step {
        /// How much higher the parent of the described node is compared to the previous node
        up: usize,
        /// The branch leading from the parent to the described node
        branch: B,
        /// Node data
        data: N,
    },
}

impl<N, B> Traversal<N, B> {
    pub fn take(self) -> N {
        match self {
            Self::Start(data) => data,
            Self::Step {
                up: _,
                branch: _,
                data,
            } => data,
        }
    }

    pub fn map<M>(self, op: impl Fn(N) -> M) -> Traversal<M, B> {
        match self {
            Self::Start(data) => Traversal::Start(op(data)),
            Self::Step { up, branch, data } => Traversal::Step {
                up,
                branch,
                data: op(data),
            },
        }
    }
    pub fn data(&self) -> &N {
        match self {
            Self::Start(data) => data,
            Self::Step {
                up: _,
                branch: _,
                data,
            } => data,
        }
    }
    pub fn branch(&self) -> Option<&B> {
        match self {
            Self::Start(_) => None,
            Self::Step {
                up: _,
                branch,
                data: _,
            } => Some(branch),
        }
    }
    pub fn up(&self) -> usize {
        match self {
            Self::Start(_) => 0,
            Self::Step {
                up,
                branch: _,
                data: _,
            } => *up,
        }
    }
    pub(crate) fn truncate_up(&mut self, cut: usize) {
        if let Self::Step { up, branch, data } = self {
            *up -= cut;
        }
    }

    pub fn is_root(&self) -> bool {
        match self {
            Traversal::Start(_) => true,
            Traversal::Step { .. } => false,
        }
    }
}

///////////////////////////////////////////////////////////////////////////////

pub trait TreeIterTarget<'a, N, B> {
    type Target;
    fn target(tree: &'a Tree<N, B>, idx: NodeIDX) -> Self::Target;
}

impl<'a, N, B> TreeIterTarget<'a, N, B> for &'a TreeNode<N, B> {
    type Target = Self;
    fn target(tree: &'a Tree<N, B>, idx: NodeIDX) -> Self::Target {
        &tree.nodes[idx]
    }
}

impl<'a, N, B> TreeIterTarget<'a, N, B> for &'a N {
    type Target = Self;
    fn target(tree: &'a Tree<N, B>, idx: NodeIDX) -> Self::Target {
        &tree.nodes[idx].value
    }
}

impl<'a, N, B> TreeIterTarget<'a, N, B> for () {
    type Target = Self;
    fn target(_: &'a Tree<N, B>, _: NodeIDX) -> Self::Target {}
}

impl<'a, N, B> TreeIterTarget<'a, N, B> for NodeIDX {
    type Target = Self;
    fn target(_: &'a Tree<N, B>, idx: NodeIDX) -> Self::Target {
        idx
    }
}

pub struct Iter<'a, N, B, T>
where
    T: TreeIterTarget<'a, N, B>,
{
    tree: &'a Tree<N, B>,
    root: Option<NodeIDX>,
    stages: Vec<HMIter<'a, B, usize>>,
    up: usize,
    _phantom: PhantomData<T>,
    //path: Path<B>,
}

impl<'a, N, B, T> Iter<'a, N, B, T>
where
    T: TreeIterTarget<'a, N, B>,
{
    pub fn new(tree: &'a Tree<N, B>) -> Self {
        Self {
            tree,
            root: tree.get_root_idx(),
            stages: vec![],
            up: 0,
            _phantom: PhantomData,
            //path: Path::new(),
        }
    }

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

impl<'a, N, B, T> Iter<'a, N, B, T>
where
    B: Clone + Eq + Hash,
    T: TreeIterTarget<'a, N, B>,
{
    pub fn new_sub_state(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,
            _phantom: PhantomData,
            //path,
        })
    }

    pub fn is_last_child(&mut self) -> bool {
        if let Some(last) = self.stages.last_mut() {
            last.peekable().peek().is_none()
        } else {
            true
        }
    }
}
//Vec<HMIterator<B, TreeNode<N, B>>>,
//);

impl<'a, N, B, T> Iterator for Iter<'a, N, B, T>
where
    T: TreeIterTarget<'a, N, B>,
{
    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.next() {
                // Go down to a child
                let node = &self.tree.nodes[*node_idx];
                self.stages.push(node.children.iter());
                //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());
            Some(Traversal::Start(T::target(self.tree, root)))
        } else {
            None
        }
    }
}

impl<'a, N, B, T> FusedIterator for Iter<'a, N, B, T> where T: TreeIterTarget<'a, N, B> {}

impl<'a, N, B> IntoIterator for &'a Tree<N, B> {
    type Item = Traversal<&'a N, &'a B>;

    type IntoIter = Iter<'a, N, B, &'a N>;

    fn into_iter(self) -> Self::IntoIter {
        Self::IntoIter::new(self)
    }
}

///////////////////////////////////////////////////////////////////////////////
//                                  IterMut                                  //
///////////////////////////////////////////////////////////////////////////////

//TODO

///////////////////////////////////////////////////////////////////////////////
//                                 IntoIter                                  //
///////////////////////////////////////////////////////////////////////////////

pub trait TreeIntoIterTarget<N, B> {
    type Target;
    fn target(node: TreeNode<N, B>) -> Self::Target;
}

impl<N, B> TreeIntoIterTarget<N, B> for N {
    type Target = Self;
    fn target(node: TreeNode<N, B>) -> Self::Target {
        node.value
    }
}

impl<N, B> TreeIntoIterTarget<N, B> for TreeNode<N, B> {
    type Target = Self;
    fn target(node: TreeNode<N, B>) -> Self::Target {
        node
    }
}

/// Tree iterator
pub struct IntoIter<N, B, T>
where
    B: Clone,
    T: TreeIntoIterTarget<N, B>,
{
    idxs: Vec<Traversal<NodeIDX, B>>,
    nodes: Vec<Option<TreeNode<N, B>>>,
    _phantom: PhantomData<T>,
}

impl<N, B, T> IntoIter<N, B, T>
where
    B: Clone,
    T: TreeIntoIterTarget<N, B>,
{
    pub fn new(tree: Tree<N, B>) -> Self {
        let mut s = Self {
            //TODO: Remove B: Clone
            idxs: Iter::<N, B, NodeIDX>::new(&tree)
                .map(|n| match n {
                    Traversal::Start(idx) => Traversal::Start(idx),
                    Traversal::Step { up, branch, data } => Traversal::Step {
                        up,
                        branch: branch.clone(),
                        data,
                    },
                })
                .collect(),
            nodes: tree.nodes.into_iter().map(Some).collect(),
            _phantom: PhantomData,
        };
        // The vector's last value will be popped on each iteration
        s.idxs.reverse();
        s
    }
}

impl<N, B, T> Iterator for IntoIter<N, B, T>
where
    B: Clone,
    T: TreeIntoIterTarget<N, B>,
{
    type Item = Traversal<T::Target, B>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.idxs.pop() {
            Some(Traversal::Start(idx)) => {
                Some(Traversal::Start(T::target(self.nodes[idx].take().unwrap())))
            }
            Some(Traversal::Step { up, branch, data }) => Some(Traversal::Step {
                up,
                branch,
                data: T::target(self.nodes[data].take().unwrap()),
            }),
            None => None,
        }
    }
}

impl<N, B> IntoIterator for Tree<N, B>
where
    B: Clone,
{
    type Item = Traversal<N, B>;

    type IntoIter = IntoIter<N, B, N>;

    fn into_iter(self) -> Self::IntoIter {
        Self::IntoIter::new(self)
    }
}