{-| Combining and composing trees of trees.  -}
module Dodge.Tree.Compose
    ( expandTree
    , singleUseAll
    , singleUseNone
    , passUntilUseAll
    , chainUses
    , useAllAtEnd
    , changeToPassDown
    , overwriteLabel
    ) where
import Dodge.Tree.Compose.Data
import Dodge.Tree.Polymorphic
import Dodge.Base
import Data.Tree

expandTree :: CompTree a -> Tree a
expandTree (Node root extChildren) = case root of
    Node (UseAll x)     _   -> Node x (map expandTree extChildren)
    Node (UseSome is x) _   -> Node x (map (expandTree . \i -> extChildren !! i) is)
    Node (UseNone _) _   -> fmap _unCompose root
    Node (UseLabel _ _) _   -> fmap _unCompose root
    Node (PassDown x)   xs  -> Node x $ map (expandTree . (`Node` extChildren)) xs
    Node (SplitDown x)  xs  -> Node x $ map expandTree $ zipWith Node xs 
            $ map (:[]) extChildren ++ repeat []

overwriteLabel :: Int -> (a -> ComposingNode a) -> SubCompTree a -> [SubCompTree a] -> SubCompTree a
overwriteLabel i f t ts = errorHead ("tried to overwriteLabel " ++ show i)
    $ updateSingleNode islabel update t
  where
    islabel (UseLabel j _) | i == j = True
    islabel _ = False
    update (Node cr _) = Node (f $ _unCompose cr) ts

passUntilUseAll :: SubCompTree a -> [SubCompTree a] -> SubCompTree a
passUntilUseAll (Node (UseAll x) _) ts' = Node (PassDown x) ts'
passUntilUseAll (Node cn ts) ts' = Node (PassDown (_unCompose cn)) $ map (`passUntilUseAll` ts') ts

chainUses :: [SubCompTree a] -> SubCompTree a
chainUses (t:s:ts) = chainUses (passUntilUseAll t [s]:ts)
chainUses [t] = t
chainUses []  = error "tried to concatenate an empty list of SubCompTrees"

singleUseAll :: a -> Tree (ComposingNode a)
singleUseAll x = Node (UseAll x) []

useAllAtEnd :: Tree a -> SubCompTree a
useAllAtEnd tree = case tree of
    Node x (t:ts)   -> Node (PassDown x) (useAllAtEnd t : map (fmap UseNone) ts)
    Node x []       -> Node (UseAll x) []

singleUseNone :: a -> SubCompTree a
singleUseNone x = Node (UseNone x) []

changeToPassDown :: ComposingNode a -> ComposingNode a
changeToPassDown x = PassDown (_unCompose x)