{-
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] (note that d is not the first child of b).
 -}
module Dodge.Layout.Tree.Polymorphic
    where
import Dodge.RandomHelp

import Data.Tree
import Control.Monad.State
import System.Random

{- |
Creates a linear tree.
Safe.
-}
treeFromPost :: [a] -> a -> Tree a
treeFromPost []     y = Node y []
treeFromPost (x:xs) y = Node x [treeFromPost xs y]

{- |
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 [] t     = t
treeFromTrunk (x:xs) t = Node x [treeFromTrunk xs t]

{- |
Applies a function to the root of a tree.
-}
applyToRoot :: (a -> a) -> Tree a -> Tree a
applyToRoot f (Node t ts) = Node (f t) ts

-- | 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 [] f t = applyToRoot f t
applyToNode (i:is) f (Node x xs) = Node x (ys ++ [applyToNode is f z] ++ zs)
  where
    (ys, z:zs) = splitAt i xs

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

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