{-
Helpers for the manipulation of rose trees.
Throughout, the _trunk_ refers to successive first children in the tree.
For example, in the tree 

>   Node a [ Node b [], Node c [Node d []] ]

the nodes in the trunk are [a,b].
 -}
module TreeHelp
    ( module Data.Tree
    , module Data.Tree.Lens
    , applyToSubtree
    , applyToSubforest
    , treeFromPost
    , treeFromTrunk
    , splitTrunk
    , applyToRandomNode
    , addToTrunk
    , inorderNumberTree
    , updateSingleNodes
    , updateRandNode
    , safeUpdateSingleNode
    , msafeUpdateSingleNode
    ) where
import RandomHelp

import Data.Maybe
import Data.Tree.Lens
import Data.Tree
import Control.Lens

{- | Creates a linear tree.
Safe.  -}
treeFromPost :: [a] -> a -> Tree a
treeFromPost xs = treeFromTrunk xs . pure

{- | Creates a tree with one trunk branch, 
input as a list, that ends in another tree.  -}
treeFromTrunk 
    :: [a]      -- ^ The trunk
    -> Tree a   -- ^ The end of the tree
    -> Tree a
treeFromTrunk = flip $ foldr f
  where
    f x t = Node x [t]

-- find use for?
---- | Consider defining this using generalised recursion patterns
--treeSize :: Tree a -> Int
--treeSize = length . flatten

{- | Applies a function to a specific node determined by a list of indices.
Unsafe (partial function).  -}
applyToNode :: [Int] -> (a -> a) -> Tree a -> Tree a
applyToNode is = applyToSubtree is . over root

{- | Applies a function to a specific subtree determined by a list of indices.
Unsafe (partial function).  -}
applyToSubtree :: [Int] -> (Tree a -> Tree a) -> Tree a -> Tree a
applyToSubtree [] f t = f t
applyToSubtree (i:is) f (Node x xs) = Node x (xs & ix i %~ applyToSubtree is f)

{- | Applies a function to a specific subforest determined by a list of indices.
Unsafe (partial function).  -}
applyToSubforest :: [Int] -> ([Tree a] -> [Tree a]) -> Tree a -> Tree a
applyToSubforest is = applyToSubtree is . over branches
--applyToSubforest [] f (Node p cs) = Node p (f cs)
--applyToSubforest (i:is) f (Node x xs) = Node x (ys ++ [applyToSubforest is f z] ++ zs)
--  where
--    (ys, z:zs) = splitAt i xs

-- do not delete: find use for
--{- |
--Applies a function to the first node along a trunk that satisfies a given property.
---}
--applyToSubTrunkBy :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> Tree a
--applyToSubTrunkBy cond f (Node x (t:ts))
--    | cond x    = f (Node x (t:ts))
--    | otherwise = Node x (applyToSubTrunkBy cond f t : ts)
--applyToSubTrunkBy _ _ t = t

-- gives the list of all updates to a single node
-- there must be a better way of doing something like this
updateSingleNodes :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> [Tree a]
updateSingleNodes f update t@(Node x ts)
    | f x       = update t : updateChildren
    | otherwise = updateChildren
  where
    updateChildren = map (Node x) (subMap (updateSingleNodes f update) ts)

mupdateSingleNodes :: Monad m => (a -> Bool) -> (Tree a -> m (Tree a)) -> Tree a -> [m (Tree a)]
mupdateSingleNodes f update t@(Node x ts)
    | f x       = update t : updateChildren
    | otherwise = updateChildren
  where
    updateChildren = map (fmap $ Node x) (msubMap (mupdateSingleNodes f update) ts) 

updateRandNode :: RandomGen g => (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> State g (Tree a)
updateRandNode t f = takeOne . updateSingleNodes t f

safeUpdateSingleNode :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> Tree a
safeUpdateSingleNode f g t = fromMaybe t $ listToMaybe $ updateSingleNodes f g t

msafeUpdateSingleNode :: Monoid b => (a -> Bool) -> (Tree a -> (b, Tree a)) -> Tree a -> (b,Tree a)
msafeUpdateSingleNode f g t = fromMaybe (mempty,t) $ listToMaybe $ mupdateSingleNodes f g t

msubMap :: Functor m => (a -> [m a]) -> [a] -> [m [a]]
msubMap f (x:xs) = (f x <&> fmap (: xs))  ++ ( (fmap (x :)) <$> msubMap f xs ) 
msubMap _ [] = []

subMap :: (a -> [a]) -> [a] -> [[a]]
subMap f (x:xs) = (f x <&> (: xs)) ++ ( (x :) <$> subMap f xs ) 
subMap _ [] = []

-- find use for?
--zipTree :: Tree a -> Tree b -> Tree (a,b)
--zipTree (Node x xs) (Node y ys) = Node (x,y) $ zipWith zipTree xs ys

{- | Makes each node into its child number, i.e. the index it has 
in the list of children of its parent.  -}
treeChildNums :: Tree a -> Tree Int
treeChildNums = setRoot 0
  where 
    setRoot :: Int -> Tree a -> Tree Int
    setRoot i (Node _ xs) = Node i (zipWith setRoot [0..] xs)

{- | Makes each node into its path, i.e. the list of indices that, 
when followed from the root, lead to the node.  -}
treePaths :: Tree a -> Tree [a]
treePaths (Node x xs) = (x :) <$> Node [] (map treePaths xs)

{- | Picks a random path in the tree.
Uniform probability that the path leads to any specific node.  -}
randomPath :: RandomGen g => Tree a -> State g [Int]
randomPath = takeOne . flatten . treePaths . treeChildNums

{- | Apply a function to the value of a node; 
the node is picked uniformly at random.  -}
applyToRandomNode :: RandomGen g => (a -> a) -> Tree a -> State g (Tree a)
applyToRandomNode f t = do
    p <- randomPath t
    return $ applyToNode p f t

{- | Add a forest to the end of a tree (along the trunk).  -}
addToTrunk :: Tree a -> [Tree a] -> Tree a
addToTrunk (Node x []) f = Node x f
addToTrunk (Node x (t:ts)) f = Node x (addToTrunk t f : ts)

{- | Find the depth of a tree along the trunk.  -}
trunkDepth :: Tree a -> Int
trunkDepth (Node _ []) = 0
trunkDepth (Node _ (x:_)) = trunkDepth x + 1

{- | Split a tree at a given point along its trunk.  -}
splitTrunkAt 
    :: Int      -- ^ Split depth
    -> Tree a -> (Tree a, [Tree a])
splitTrunkAt 0 (Node x xs) = (Node x [],xs)
splitTrunkAt i (Node y (x:xs)) = 
    let (t, ts) = splitTrunkAt (i-1) x
    in (Node y (t : xs) , ts)
splitTrunkAt _ (Node _ []) = error "Trying to split to short a trunk"

{- | Split a tree at a random point along its trunk.  -}
splitTrunk :: RandomGen g => Tree a -> State g (Tree a, [Tree a])
splitTrunk t = do
    i <- state $ randomR (0, trunkDepth t)
    return $ splitTrunkAt i t

-- untested
inorderNumberTree :: Tree a -> Tree (a,Int)
inorderNumberTree = fst . f 0
  where
    f i (Node x ts) =
        let (ts',i') = g (i+1) ts
        in (Node (x,i) ts', i')
    g i (t:ts) = 
        let (t',i') = f i t
            (ts',i'') = g i' ts
        in (t': ts', i'')
    g i [] = ([], i)