import Tries
import Maybe

data BinTree = Empty | Branch BinTree String BinTree
data MapBinTree a
  = TrieBinTree (Maybe a) (Maybe (MapBinTree (MapStr (MapBinTree a))))

mapStrToList :: MapStr a -> [(String,a)]
mapStrToList (TrieStr nil cons)
  = [ ("",x)   | x <- maybeToList nil ]
 ++ [ (c:cs,x) | (c,sm) <- cons, (cs,x) <- mapStrToList sm ]

checkStr :: (a -> Bool) -> MapStr a -> Bool
checkStr _   (TrieStr Nothing []) = False
checkStr chk (TrieStr _ cons) = all (checkStr chk . snd) cons

combineStr :: (a -> a -> a) -> MapStr a -> MapStr a -> MapStr a
combineStr cmb (TrieStr mn1 c1) (TrieStr mn2 c2)
  = TrieStr (maybe mn2 (Just . \n1 -> maybe n1 (cmb n1) mn2) mn1) (c1++c2)

mapTreeToList :: MapTree a -> [(Tree,a)]
mapTreeToList (TrieTree leaf node)
  = [ (Leaf s  , x) | sm <- maybeToList leaf, (s,x) <- mapStrToList sm ]
 ++ [ (Node l r, x) | lm <- maybeToList node
                    , (l,rm) <- mapTreeToList lm
                    , (r,x)  <- mapTreeToList rm ]

mapTreeIsTidy :: MapTree a -> Bool
mapTreeIsTidy (TrieTree Nothing Nothing) = True
mapTreeIsTidy mt = checkTree (const True) mt

checkTree :: (a -> Bool) -> MapTree a -> Bool
checkTree _   (TrieTree Nothing Nothing) = False
checkTree chk (TrieTree leaf node)
  = maybe True (checkStr chk) leaf
 && maybe True (checkTree (checkTree (const True))) node

mapTreeUnion :: MapTree a -> MapTree a -> MapTree a
mapTreeUnion = combineTree const

combineTree :: (a -> a -> a) -> MapTree a -> MapTree a -> MapTree a
combineTree cmb (TrieTree ml1 mn1) (TrieTree ml2 mn2)
  = TrieTree
     (maybe ml2 (\l1 -> maybe ml1 (tidyStr . combineStr cmb l1) ml2) ml1)
     (maybe mn2
            (\n1 -> maybe mn1
                          (tidyTree . combineTree (combineTree cmb) n1)
                          mn2)
            mn1)