use IntoIterator;
use Iterator;

pub mod tree_format;

extern crate boolinator;
use boolinator::Boolinator;

use std::{
  collections::VecDeque,
  fmt::{Debug, Display},
};

use crate::tree_format::TreeFormat;

#[derive(Debug, Clone, Eq)]
pub enum RoseTree<V> {
  /// A simple (and almost certainly terribly inefficient) Rose Tree
  /// implementation.

  /// A dead end node, single value.
  Leaf(V),
  /// A branch has a head value, and then all its child values.
  Branch(V, Vec<RoseTree<V>>),
}

impl<V> From<V> for RoseTree<V> {
  /// Convenience function to make wrapping anything in a leaf nicer.
  fn from(val: V) -> Self {
    RoseTree::Leaf(val)
  }
}

impl<V: Display> Display for RoseTree<V> {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
    write!(f, "{}", self.fmt_tree())
  }
}

impl<V, I: IntoIterator<Item = V>> From<(V, I)> for RoseTree<V> {
  /// Convenience to make wrapping anything in a Branch nicer(I know, shocking).
  fn from((head, child_collection): (V, I)) -> RoseTree<V> {
    let children = child_collection.into_iter().map(Into::into).collect();
    RoseTree::Branch(head, children)
  }
}

impl<V: PartialEq> PartialEq for RoseTree<V> {
  /// It'd be really nice to be able to check equality of variants
  /// independently of the values within the variants, but I have no idea how
  /// or if that's possible.
  fn eq(&self, other: &Self) -> bool {
    match self {
      RoseTree::Leaf(vl) => match other {
        RoseTree::Leaf(vr) => vl == vr,
        _ => false,
      },
      RoseTree::Branch(head, children) => match other {
        RoseTree::Branch(other_head, other_children) => {
          if head != other_head {
            return false;
          };
          if children.len() != other_children.len() {
            return false;
          };
          children
            .iter()
            .zip(other_children.iter())
            .map(|(child, other_child)| child == other_child)
            .fold(true, |accum, is_equal| accum && is_equal)
        }
        _ => false,
      },
    }
  }
}

impl<V> RoseTree<V> {
  /// Constructors, because I don't know how to make `From` do this for me.
  pub fn leaf(val: V) -> Self {
    RoseTree::Leaf(val)
  }

  pub fn branch<I: Iterator<Item = V>>(val: V, children: I) -> Self {
    let new_children = children.into_iter().map(|c| c.into()).collect();
    Self::Branch(val, new_children)
  }

  /// Creates a tree of the same structure as `self` where every node is a
  /// reference to the corresponding node in the owning tree.
  pub fn ref_tree(&self) -> RoseTree<&V> {
    match self {
      RoseTree::Leaf(val) => RoseTree::Leaf(val),
      RoseTree::Branch(head, children) => {
        let ref_children = children.iter().map(RoseTree::ref_tree).collect();
        RoseTree::Branch(head, ref_children)
      }
    }
  }

  /// These three check if a RoseTree is of a given variant.

  pub fn is_leaf(&self) -> bool {
    match self {
      RoseTree::Leaf(_) => true,
      _ => false,
    }
  }

  pub fn is_branch(&self) -> bool {
    match self {
      RoseTree::Branch(_, _) => true,
      _ => false,
    }
  }

  /// Calculates the number of nodes in the tree.
  pub fn size(&self) -> usize {
    match self {
      RoseTree::Leaf(_) => 1,
      RoseTree::Branch(_, children) => 1 + children.iter().map(|c| c.size()).sum::<usize>(),
    }
  }

  /// Retrieves a reference to the head of the tree.
  pub fn head(&self) -> Option<&V> {
    match self {
      RoseTree::Leaf(val) => Some(val),
      RoseTree::Branch(head, _) => Some(head),
    }
  }

  pub fn head_mut(&mut self) -> Option<&mut V> {
    match self {
      RoseTree::Leaf(val) => Some(val),
      RoseTree::Branch(head, _) => Some(head),
    }
  }

  /// Gets the children of a branch directly, if it is a branch.
  pub fn children(&self) -> Option<&Vec<RoseTree<V>>> {
    match self {
      RoseTree::Branch(_, chs) => Some(chs),
      _ => None,
    }
  }

  pub fn children_mut(&mut self) -> Option<&mut Vec<RoseTree<V>>> {
    match self {
      RoseTree::Branch(_, chs) => Some(chs),
      _ => None,
    }
  }

  /// Swaps out the `head-most` value of the tree with the provided one.
  pub fn swap_head(mut self, new_head: V) -> Self {
    match self {
      RoseTree::Leaf(_) => RoseTree::Leaf(new_head),
      RoseTree::Branch(ref mut head, _) => {
        *head = new_head;
        self
      }
    }
  }

  /// Adds a child to specifically a Branch node.
  pub fn add_child(self, child: V) -> Self {
    match self {
      RoseTree::Branch(head, mut children) => {
        children.push(child.into());
        RoseTree::Branch(head, children)
      }
      _ => self,
    }
  }

  // Maps a function over just the head of the tree.
  pub fn map_head(self, head_fn: fn(V) -> V) -> Self {
    match self {
      RoseTree::Leaf(val) => RoseTree::Leaf(head_fn(val)),
      RoseTree::Branch(head, children) => RoseTree::Branch(head_fn(head), children),
    }
  }

  /// Map over only the leaves of the tree.
  pub fn map_leaves(self, child_fn: fn(V) -> V) -> Self {
    match self {
      RoseTree::Leaf(_) => self,
      RoseTree::Branch(head, children) => RoseTree::Branch(
        head,
        children
          .into_iter()
          .map(|child| child.map_leaves(child_fn))
          .collect(),
      ),
    }
  }

  /// Map over the whole tree.
  pub fn map<U>(&self, func: fn(&V) -> U) -> RoseTree<U> {
    match self {
      RoseTree::Leaf(val) => RoseTree::Leaf(func(val)),
      RoseTree::Branch(head, children) => RoseTree::Branch(
        func(head),
        children.iter().map(|child| child.map(func)).collect(),
      ),
    }
  }

  pub fn map_into<U>(self, func: fn(V) -> U) -> RoseTree<U> {
    match self {
      RoseTree::Leaf(val) => RoseTree::Leaf(func(val)),
      RoseTree::Branch(head, children) => RoseTree::Branch(
        func(head),
        children
          .into_iter()
          .map(|child| child.map_into(func))
          .collect(),
      ),
    }
  }

  /// Maps over the tree, where each child of the same parent receives the same
  /// accumulator value as a function argument.
  /// Add more variants of this for different argument mutabilities maybe?
  pub fn depth_map<A: Clone, U>(&self, initial: A, func: fn(A, &V) -> (A, U)) -> RoseTree<U> {
    match self {
      RoseTree::Leaf(val) => {
        let (_, new_val) = func(initial, val);
        RoseTree::Leaf(new_val)
      }
      RoseTree::Branch(head, children) => {
        let (accum, new_head) = func(initial, head);
        let mut new_children = Vec::with_capacity(children.len());
        for child in children {
          new_children.push(child.depth_map(accum.clone(), func))
        }
        RoseTree::Branch(new_head, new_children)
      }
    }
  }

  /// Removes all values from the tree which don't satisfy the predicate. In the
  /// case of branches, if the head of the branch does not satisfy the
  /// predicate, both it and all its children are removed.
  pub fn filter<F: Fn(&V) -> bool>(self, func: F) -> Option<Self> {
    match self {
      RoseTree::Leaf(val) => func(&val).as_some(RoseTree::Leaf(val)),
      RoseTree::Branch(head, children) => func(&head).as_some_from(|| {
        RoseTree::Branch(
          head,
          children
            .into_iter()
            .filter_map(|c| c.filter(&func))
            .collect(),
        )
      }),
    }
  }

  // I'll let you guess on this one. I believe in you.
  pub fn preorder_iter<'a>(&'a self) -> PreorderRoseIter<V> {
    let mut stack = VecDeque::new();
    stack.push_back(self);
    PreorderRoseIter { stack }
  }
}

/// A depth-first iterator over references to the tree values. No points for
/// guessing.
pub struct PreorderRoseIter<'a, T> {
  stack: VecDeque<&'a RoseTree<T>>,
}

impl<'a, T: Debug> Iterator for PreorderRoseIter<'a, T> {
  type Item = &'a T;

  fn next(&mut self) -> Option<Self::Item> {
    let next_node = self.stack.pop_front()?;
    match next_node {
      RoseTree::Leaf(ref val) => Some(val),
      RoseTree::Branch(ref head, ref children) => {
        for child in children.iter().rev() {
          self.stack.push_front(child);
        }
        Some(head)
      }
    }
  }
}

#[cfg(test)]
mod test {
  use crate::RoseTree;

  fn test_tree() -> RoseTree<usize> {
    RoseTree::Branch(
      1,
      vec![
        RoseTree::Branch(
          2,
          vec![RoseTree::Leaf(3), RoseTree::Leaf(4), RoseTree::Leaf(5)],
        ),
        RoseTree::Branch(
          6,
          vec![RoseTree::Leaf(7), RoseTree::Leaf(8), RoseTree::Leaf(9)],
        ),
        RoseTree::Branch(
          10,
          vec![RoseTree::Leaf(11), RoseTree::Leaf(12), RoseTree::Leaf(13)],
        ),
      ],
    )
  }

  #[test]
  fn test_preorder_actually_preordered() {
    let expected: Vec<usize> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13];
    let actual: Vec<usize> = test_tree().preorder_iter().map(|v| *v).collect();
    assert_eq!(expected, actual);
  }
}