{-|
Combining and composing trees.
-}
module Dodge.Layout.Tree
    where
import Data.Tree
import Control.Monad.State
import System.Random
-- | 'Left' elements get new children, 'Right' elements inherit the children from the
-- mapped over node
expandTreeBy :: (a -> Tree (Either b b)) -> Tree a -> Tree b
expandTreeBy f (Node x []) = fmap removeEither (f x)
    where removeEither (Left y)  = y
          removeEither (Right y) = y
expandTreeBy f (Node x xs) = appendAndRemove $ f x
    where appendAndRemove (Node (Left y)  ys) = Node y (map appendAndRemove ys)
          appendAndRemove (Node (Right y) _ ) = Node y (map (expandTreeBy f) xs)

expandTreeRand :: RandomGen g => 
        (a -> State g (Tree (Either b b))) -> Tree a -> State g (Tree b)
expandTreeRand f (Node x []) = fmap (fmap removeEither) (f x)
    where removeEither (Left y)  = y
          removeEither (Right y) = y
expandTreeRand f (Node x xs) = do
    root <- f x
    branches <- sequence $ map (expandTreeRand f) xs
    return (appendAndRemove branches root)
        where appendAndRemove :: [Tree a] -> Tree (Either a a) -> Tree a
              appendAndRemove bran (Node (Left y) ys) = Node y (map (appendAndRemove bran) ys)
              appendAndRemove bran (Node (Right y) _) = Node y bran

collapseTree :: Tree (Tree (Either a a)) -> Tree (Either a a)
collapseTree = g . expandTreeBy f
    where f (Node (Right x) xs) = Node (Right (Right x)) $ map f xs
          f (Node (Left x) xs) = Node (Left (Left x)) $ map f xs
          g (Node (Right x) []) = Node (Right x) []
          g (Node (Right x) xs) = Node (Left x) $ map g xs
          g (Node y ys) = Node y $ map g ys

treePost :: [a] -> a -> Tree a
treePost []     y = Node y []
treePost (x:xs) y = Node x [treePost xs y]

treeTrunk :: [a] -> Tree a -> Tree a
treeTrunk [] t     = t
treeTrunk (x:xs) t = Node x [treeTrunk xs t]

applyToRoot :: (a -> a) -> Tree a -> Tree a
applyToRoot f (Node x xs) = Node (f x) xs