{-# LANGUAGE TupleSections #-}
{-| Combining and composing trees of trees.  -}
module Dodge.Tree.Compose
    ( tToBTree
    , overwriteLabel
    , composeTree
    , composeTree'
    , shiftChildren
    , attachTree
    , composeTreeRand
    , composeAndLog
    , toOnward
    , attachOnward
--    , cmpToMT
--    , compTree
    , module Dodge.Tree.Compose.Data
    ) where
import Dodge.Data
import Dodge.Tree.Compose.Data
--import Dodge.RoomCluster.Data
import TreeHelp
--import Dodge.Base
import LensHelp

import Data.Bifunctor
import Control.Monad.State
import Data.Foldable
import System.Random
import Data.Maybe

--expandTree :: CompTree a -> Tree a
--expandTree (Node rt extChildren) = case rt 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 rt
--    Node (useLabel _ _) _   -> fmap _unCompose rt
--    Node ( x)   xs  -> Node x $ map (expandTree . (`Node` extChildren)) xs
--    Node (SplitDown x)  xs  -> Node x $ map expandTree $ zipWith Node xs 
--            $ map (:[]) extChildren ++ repeat []

--applyToCompRoot :: (a -> a) -> Tree a -> Tree a
--applyToCompRoot f = over root (unCompose %~ f)

overwriteLabel :: Int -> Tree Room -> [Tree Room] -> Tree Room
overwriteLabel i t ts = safeUpdateSingleNode 
    (\rm -> LabelCluster i `elem` (rm ^?! rmClusterStatus . csLinks))
    ((branches .~ ts) . (root . rmClusterStatus . csLinks . at (LabelCluster i) .~ Nothing))
    t

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

--passUntiluseAll (Node (useAll x) _) ts' = Node ( x) ts'
--passUntiluseAll (Node cn ts) ts' = Node ( (_unCompose cn)) $ map (`passUntiluseAll` ts') ts

--chainUses :: [Tree a] -> Tree a
--chainUses (t:s:ts) = chainUses (passUntiluseAll t [s]:ts)
--chainUses [t] = t
--chainUses []  = error "tried to concatenate an empty list of Trees"

--pure :: a -> Tree (ComposingNode a)
--pure x = Node (useAll x) []

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

--pure :: a -> Tree a
--pure x = Node (UseNone x) []

--changeTo :: ComposingNode a -> ComposingNode a
--changeTo x =  (_unCompose x)

composeTreeRand :: RandomGen g => Tree (a -> Maybe a, Tree a) -> State g (Tree a)
composeTreeRand (Node (_,t) ts) = attachListR t ts

attachListR :: RandomGen g => Tree a -> [Tree (a -> Maybe a, Tree a)] -> State g (Tree a)
attachListR t xs = shiftChildren <$> foldrM newfRand (fmap ([],) t) xs

newfRand :: RandomGen g => Tree (a -> Maybe a,Tree a) -> Tree ([Tree a],a) 
    -> State g (Tree ([Tree a],a))
newfRand t = updateRandNode (isJust . upf . snd) ( (root . _2 %~ g) . (root . _1 .:~ composeTree t ) )
  where
    upf = t ^?! root . _1
    g x = fromMaybe x (upf x)

composeTree :: Tree (a -> Maybe a, Tree a) -> Tree a
composeTree (Node (_,t) ts) = attachList t ts

composeTree' :: Monoid b => Tree (a -> Maybe (b, a), Tree a) -> (b, Tree a)
composeTree' (Node (_,t) ts) = attachList' t ts

attachList' :: Monoid b => Tree a -> [Tree (a -> Maybe (b, a), Tree a)] -> (b,Tree a)
attachList' t xs = shiftChildren <$> foldrM newf' (fmap ([],) t) xs

newf' :: Monoid b => Tree (a -> Maybe (b, a),Tree a) -> Tree ([Tree a],a) -> (b , Tree ([Tree a],a))
newf' t = first (s<>) . msafeUpdateSingleNode (isJust . upf . snd) upt
  where
    (s,t') = composeTree' t
    upf = t ^?! root . _1
    upt = (root . _2 %%~ g) . (root . _1 .:~ t')
    g x = fromMaybe (mempty,x) (upf x)
--    mtf = undefined
--     ( (root . _2 %~ g) . (root . _1 .:~ composeTree t ) )

attachList :: Tree a -> [Tree (a -> Maybe a, Tree a)] -> Tree a
attachList t xs = shiftChildren $ foldr newf (fmap ([],) t) xs

newf :: Tree (a -> Maybe a,Tree a) -> Tree ([Tree a],a) -> Tree ([Tree a],a)
newf t = safeUpdateSingleNode (isJust . upf . snd) ( (root . _2 %~ g) . (root . _1 .:~ composeTree t ) )
  where
    upf = t ^?! root . _1
    g x = fromMaybe x (upf x)

attachTree :: (a -> Maybe a) -> Tree a -> Tree ([Tree a],a) -> Tree ([Tree a],a)
attachTree upf t = safeUpdateSingleNode (isJust . upf . snd) ( (root . _2 %~ g) . (root . _1 .:~ t ) )
  where
    g x = fromMaybe x (upf x)

shiftChildren :: Tree ([Tree a],a) -> Tree a
shiftChildren (Node (ts,x) ts') = Node x (ts ++ map shiftChildren ts')

---- hideous, but it should work
---- the idea is that the fst argument selects (and edits) a node to which all the
---- children are attached
---- the tree is thus built up recursively
--compTree :: Tree (a -> Maybe a, Tree (a,b)) -> Tree b
--compTree (Node (_,t) ch) = attachChildren t $ map (\t' -> (fst $ rootLabel t',compTree t')) ch
--
--chooseAttachNodes :: (a -> Maybe a,c) -> Tree (a,b,[c]) -> [Tree (a,b,[c])]
--chooseAttachNodes (t,z) (Node (x,y,zs) ch) = case t x of
--    Just x' -> Node (x',y,(z:zs)) ch : concatMap (chooseAttachNodes (t,z)) ch
--    Nothing -> concatMap (chooseAttachNodes (t,z)) ch
--
--mAttachToNode :: (a -> Maybe a,c) -> Tree (a,b,[c]) -> Tree (a,b,[c])
--mAttachToNode p t = fromMaybe t $ listToMaybe $ chooseAttachNodes p t
--
--moveToChildren :: Tree (a,b,[Tree b]) -> Tree b
--moveToChildren (Node (_,x,ts) ts') = Node x (map moveToChildren ts' ++ ts)
--
--attachChildren :: Tree (a,b) -> [(a->Maybe a,Tree b)] -> Tree b
--attachChildren t = moveToChildren . foldr mAttachToNode t'
--  where
--    t' = fmap (\(x,y) -> (x,y,[])) t

logMB :: MetaBranch a -> Tree String
logMB (MBranch str _ mt) = root ++.~ (str ++ ":") $ logMT mt

logMT :: MetaTree a -> Tree String
logMT mt = case mt of
    MTree str _ brs -> Node ("M:"++str) $ map logMB brs
    BTree str _ brs -> Node str $ map logMT brs

composeAndSubLog :: Monad m => [Int] -> (String -> m ()) -> MetaTree a -> m (Tree a)
composeAndSubLog is dolog mt = case mt of
    MTree{} -> do
        dolog ("SUBTREE"++show is++mRootLabel mt++"\n")
        rt <- composeAndLog (0:is) dolog $ _mtTree mt
        brs <- mapM (\(i,mb) -> (_mbAttach mb,) <$> composeAndSubLog (is & ix 0 +~ 1) dolog (_mbTree mb)) 
            $ zip [0..] $ _mtBranches mt
        return $ attachList rt $ map pure brs
    bt -> do
        dolog ("SUBTREE:"++show is++mRootLabel mt++"\n")
        dolog $ drawTree $ logMT mt
        brs <- mapM (composeAndSubLog is dolog) (_btBranches bt)
        return $ Node (_btValue bt) brs

composeAndLog :: Monad m => [Int] -> (String -> m ()) -> MetaTree a -> m (Tree a)
composeAndLog is dolog mt = case mt of
    MTree{} -> do
        dolog $ drawTree $ logMT mt
        dolog ("SUBTREE"++show is++mRootLabel mt++"\n")
        rt <- composeAndLog (0:is) dolog $ _mtTree mt
        brs <- mapM (\(i,mb) -> (_mbAttach mb,) <$> composeAndSubLog (is & ix 0 +~ 1) dolog (_mbTree mb)) 
            $ zip [0..] $ _mtBranches mt
        return $ attachList rt $ map pure brs
    BTree{} -> do
        dolog ("SUBTREE:"++show is++mRootLabel mt++"\n")
        dolog $ drawTree $ logMT mt
        brs <- mapM (composeAndSubLog is dolog) $ _btBranches mt
        return $ Node (_btValue mt) brs

tToBTree :: Tree Room -> MetaTree Room
tToBTree (Node r ts) = BTree (_rmName r) r $ map tToBTree ts

attachOnward :: MetaTree Room -> MetaTree Room -> MetaTree Room
attachOnward t1 t2 = case t1 of
    MTree {} -> t1 & mtBranches .:~ MBranch "Onward" toOnward t2
    BTree {} -> MTree (mRootLabel t1 ++ "-" ++ mRootLabel t2) t1 [MBranch "Onward" toOnward t2]

toOnward :: Room -> Maybe Room
toOnward rm
    | OnwardCluster `elem` rm ^?! rmClusterStatus . csLinks 
        = Just (rm & rmClusterStatus . csLinks . at OnwardCluster .~ Nothing)
    | otherwise = Nothing