-- | Tree-related utilities.
module Data.Tree.Util where

import           Data.Bifoldable
import           Data.Bifunctor
import           Data.Bitraversable
import           Control.Lens
import           Control.Monad ((>=>))
import qualified Data.List as List
import qualified Data.List.NonEmpty as NonEmpty
import           Data.List.NonEmpty (NonEmpty(..))
import           Data.Maybe (listToMaybe,maybeToList)
import           Data.Tree

--------------------------------------------------------------------------------

-- $setup
-- >>> :{
-- let myTree = Node 0 [ Node 1 []
--                     , Node 2 []
--                     , Node 3 [ Node 4 [] ]
--                     ]
-- :}


--------------------------------------------------------------------------------

-- | Nodes in a tree are typically either an internal node or a leaf node
data TreeNode v a = InternalNode v | LeafNode a deriving (Show,Eq)

instance Bifunctor TreeNode where
  bimap = bimapDefault
instance Bifoldable TreeNode where
  bifoldMap = bifoldMapDefault
instance Bitraversable TreeNode where
  bitraverse f g = \case
    InternalNode v -> InternalNode <$> f v
    LeafNode l     -> LeafNode     <$> g l

-- | A TreeNode is isomorphic to Either
_TreeNodeEither :: Iso' (TreeNode v p) (Either v p)
_TreeNodeEither = iso tne etn
  where
    tne = \case
      InternalNode v -> Left v
      LeafNode l     -> Right l
    etn = either InternalNode LeafNode

--------------------------------------------------------------------------------
-- * Zipper on rose trees

-- | Zipper for rose trees
data Zipper a = Zipper { focus      :: Tree a
                       , ancestors  :: [([Tree a], a, [Tree a])] -- left siblings in reverse order
                       }
              deriving (Show,Eq)

-- | Create a new zipper focussiong on the root.
root :: Tree a -> Zipper a
root = flip Zipper []

-- | Move the focus to the parent of this node.
up               :: Zipper a -> Maybe (Zipper a)
up (Zipper t as) = case as of
                     []              -> Nothing
                     ((ls,p,rs):as') -> Just $ Zipper (Node p (reverse ls <> [t] <> rs)) as'

-- | Move the focus to the first child of this node.
--
-- >>> firstChild $ root myTree
-- Just (Zipper {focus = Node {rootLabel = 1, subForest = []}, ancestors = [([],0,[Node {rootLabel = 2, subForest = []},Node {rootLabel = 3, subForest = [Node {rootLabel = 4, subForest = []}]}])]})
firstChild                          :: Zipper a -> Maybe (Zipper a)
firstChild (Zipper (Node x chs) as) = case chs of
                                        []       -> Nothing
                                        (c:chs') -> Just $ Zipper c (([],x,chs'):as)

-- | Move the focus to the next sibling of this node
--
-- >>> (firstChild $ root myTree) >>= nextSibling
-- Just (Zipper {focus = Node {rootLabel = 2, subForest = []}, ancestors = [([Node {rootLabel = 1, subForest = []}],0,[Node {rootLabel = 3, subForest = [Node {rootLabel = 4, subForest = []}]}])]})
nextSibling               :: Zipper a -> Maybe (Zipper a)
nextSibling (Zipper t as) = case as of
                              []                -> Nothing -- no parent
                              ((_,_,[]):_)      -> Nothing -- no next sibling
                              ((ls,p,r:rs):as') -> Just $ Zipper r ((t:ls,p,rs):as')

-- | Move the focus to the next sibling of this node
prevSibling               :: Zipper a -> Maybe (Zipper a)
prevSibling (Zipper t as) = case as of
                              []                -> Nothing -- no parent
                              (([],_,_):_)      -> Nothing -- no prev sibling
                              ((l:ls,p,rs):as') -> Just $ Zipper l ((ls,p,t:rs):as')

-- | Given a zipper that focussses on some subtree t, construct a list with
-- zippers that focus on each child.
allChildren :: Zipper a -> [Zipper a]
allChildren = List.unfoldr ((\ch -> (ch, nextSibling ch)) <$>) . firstChild

-- | Given a zipper that focussses on some subtree t, construct a list with
-- zippers that focus on each of the nodes in the subtree of t.
allTrees   :: Zipper a -> [Zipper a]
allTrees r = r : concatMap allTrees (allChildren r)

-- | Creates a new tree from the zipper that thas the current node as root. The
-- ancestorTree (if there is any) forms the first child in this new root.
unZipperLocal                          :: Zipper a -> Tree a
unZipperLocal (Zipper (Node x chs) as) = Node x (maybeToList (constructTree as) <> chs)

-- | Constructs a tree from the list of ancestors (if there are any)
constructTree :: [([Tree a],a,[Tree a])] -> Maybe (Tree a)
constructTree = listToMaybe
              . foldr (\(ls,p,rs) tas -> [Node p (tas <> reverse ls <> rs)]) []


--------------------------------------------------------------------------------

-- | Given a predicate on an element, find a node that matches the predicate, and turn that
-- node into the root of the tree.
--
-- running time: \(O(nT)\) where \(n\) is the size of the tree, and \(T\) is
-- the time to evaluate a predicate.
--
-- >>> findEvert (== 4) myTree
-- Just (Node {rootLabel = 4, subForest = [Node {rootLabel = 3, subForest = [Node {rootLabel = 0, subForest = [Node {rootLabel = 1, subForest = []},Node {rootLabel = 2, subForest = []}]}]}]})
-- >>> findEvert (== 5) myTree
-- Nothing
findEvert   :: (a -> Bool) -> Tree a -> Maybe (Tree a)
findEvert p = findEvert' (p . rootLabel)

-- | Given a predicate matching on a subtree, find a node that matches the predicate, and turn that
-- node into the root of the tree.
--
-- running time: \(O(nT(n))\) where \(n\) is the size of the tree, and \(T(m)\) is
-- the time to evaluate a predicate on a subtree of size \(m\).
findEvert'   :: (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
findEvert' p = fmap unZipperLocal . List.find (p . focus) . allTrees . root

-- | Function to extract a path between a start node and an end node (if such a
--path exists). If there are multiple paths, no guarantees are given about
--which one is returned.
--
-- running time: \(O(n(T_p+T_s)\), where \(n\) is the size of the tree, and
-- \(T_p\) and \(T_s\) are the times it takes to evaluate the @isStartingNode@
-- and @isEndingNode@ predicates.
--
--
-- >>> findPath (== 1) (==4) myTree
-- Just [1,0,3,4]
-- >>>  findPath (== 1) (==2) myTree
-- Just [1,0,2]
-- >>>  findPath (== 1) (==1) myTree
-- Just [1]
-- >>>  findPath (== 1) (==2) myTree
-- Just [1,0,2]
-- >>>  findPath (== 4) (==2) myTree
-- Just [4,3,0,2]
findPath               :: (a -> Bool) -- ^ is this node a starting node
                          -> (a -> Bool) -- ^ is this node an ending node
                          -> Tree a -> Maybe [a]
findPath isStart isEnd = findEvert isStart >=> findNode isEnd

-- | Given a predicate on a, find (the path to) a node that satisfies the predicate.
--
-- >>> findNode (== 4) myTree
-- Just [0,3,4]
findNode   :: (a -> Bool) -> Tree a -> Maybe [a]
findNode p = listToMaybe . findNodes (p . rootLabel)

-- | Find all paths to nodes that satisfy the predicate
--
-- running time: \(O(nT(n))\) where \(n\) is the size of the tree, and \(T(m)\) is
-- the time to evaluate a predicate on a subtree of size \(m\).
--
-- >>> findNodes ((< 4) . rootLabel) myTree
-- [[0],[0,1],[0,2],[0,3]]
-- >>> findNodes (even . rootLabel) myTree
-- [[0],[0,2],[0,3,4]]
-- >>> let size = length in findNodes ((> 1) . size) myTree
-- [[0],[0,3]]
findNodes   :: (Tree a -> Bool) -> Tree a -> [[a]]
findNodes p = go
  where
    go t = let mh = [ [] | p t ] -- [[]] iff 'p t'
           in map (rootLabel t:) $ mh <> concatMap go (children t)


-- | BFS Traversal of the rose tree that decomposes it into levels.
--
-- running time: \(O(n)\)
levels :: Tree a -> NonEmpty (NonEmpty a)
levels = go1 . (:| [])
  where
    go0   :: [Tree a] -> [NonEmpty a]
    go0 q = case NonEmpty.nonEmpty q of
              Nothing -> []
              Just q1 -> NonEmpty.toList $ go1 q1
    {-# INLINE go0 #-}

    -- all work essentially happens here: given a bunch of trees whose
    -- root elements all have the same level, extract the values
    -- stored at these root nodes, collect all children in a big list,
    -- and explore those recursively.
    go1    :: NonEmpty (Tree a) -> NonEmpty (NonEmpty a)
    go1 qs = fmap root' qs :| go0 (concatMap children' qs)
    {-# INLINE go1 #-}

    root'     (Node x _)   = x
    children' (Node _ chs) = chs