{-
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].

Partial functions are annotated as such.
 -}
module TreeHelp (
    module Data.Tree,
    module Data.Tree.Lens,
    applyToSubtree,
    applyToSubforest,
    treeFromPost,
    treePost,
    treeFromTrunk,
    splitTrunk,
    applyToRandomNode,
    addToTrunk,
    inorderNumberTree,
    updateSingleNodes,
    updateAllNodes,
    updateRandNode,
    safeUpdateSingleNode,
    numTraversable,
) where

import Control.Lens
import Data.Maybe
import Data.Traversable
import Data.Tree
import Data.Tree.Lens
import RandomHelp

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

{- | Creates a linear tree from a list.
Partial function.
-}
treePost :: [a] -> Tree a
treePost xs = treeFromPost (init xs) (last xs)

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

{- | Applies a function to a specific node determined by a list of indices.
Partial.
-}
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.
Partial.
-}
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.
Partial.
-}
applyToSubforest :: [Int] -> ([Tree a] -> [Tree a]) -> Tree a -> Tree a
applyToSubforest is = applyToSubtree is . over branches

{- | Transforms a tree everywhere a property is satisfied.
Can perform multiple transformations in different branches.
However, does not perform any further transformation within a transformed subtree.
Thus this always terminate on finite trees.
-}
updateAllNodes :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> Tree a
updateAllNodes f update t@(Node x ts)
    | f x = update t
    | otherwise = updateChildren
  where
    updateChildren = Node x (map (updateAllNodes f update) ts)

{- | For each node within the input tree that satisfies a property,
produce a tree with the subtree starting at that node transformed by a function.
Produces an empty list if no nodes satisfy the property.
Depth first.
-}
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)

{- | Chooses a random node that satisfies a property,
produces a tree with the subtree starting at that node transformed by a function.
Partial, undefined if no node satisfies the property.
-}
updateRandNode :: RandomGen g => (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> State g (Tree a)
updateRandNode t f = takeOne . updateSingleNodes t f

{- | Finds the (depth) first node that satisfies a property,
produces a tree with the subtree starting at that node transformed by a function.
Produces the original tree if no nodes satisfy the property.
-}
safeUpdateSingleNode :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> Tree a
safeUpdateSingleNode f g t = fromMaybe t $ listToMaybe $ updateSingleNodes f g t

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 ::
    -- | Split depth
    Int ->
    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

numTraversable :: Traversable t => t a -> t (a, Int)
numTraversable = snd . mapAccumL f 0
  where
    f i x = (i + 1, (x, i))

-- untested
inorderNumberTree :: Tree a -> Tree (a, Int)
inorderNumberTree = numTraversable