{-# LANGUAGE BangPatterns, ViewPatterns #-}
module TreeTransform
( convertToSuperTree
, getPropertyMap
, innerToLeaves
, filterExclusiveTree
, addUniqueNodeIDs
) where
import Control.Monad.State
import Data.Function (on)
import Data.Tree
import qualified Data.Foldable as F
import qualified Data.Map.Strict as Map
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import qualified Data.Text as T
import Math.TreeFun.Types
import Math.TreeFun.Tree
import TextShow (showt)
import Types
convertToSuperTree :: Tree NodeLabel -> Tree (SuperNode NodeLabel)
convertToSuperTree = toSuperNodeTree SuperRoot
innerToLeaves :: Tree NodeLabel -> Tree NodeLabel
innerToLeaves n@(Node { subForest = [] }) = n
innerToLeaves n@( Node { rootLabel = NodeLabel { nodeLabels = (Seq.null -> True)
}
}
) = n { subForest = map innerToLeaves . subForest $ n }
innerToLeaves n@( Node { rootLabel = NodeLabel { nodeID = x }
, subForest = xs
}
) =
Node { rootLabel = NodeLabel {nodeID = T.cons 'S' x, nodeLabels = Seq.empty}
, subForest = (n { subForest = [] }) : map innerToLeaves xs
}
getPropertyMap :: Tree (SuperNode NodeLabel) -> PropertyMap NodeLabel Label
getPropertyMap = Map.fromList
. map (\ !x -> (myRootLabel x, nodeLabels . myRootLabel $ x))
. filter (not . Seq.null . nodeLabels . myRootLabel)
. leaves
filterExclusiveTree :: Exclusivity
-> Tree NodeLabel
-> Tree NodeLabel
filterExclusiveTree ex n =
n { rootLabel = (rootLabel n) { nodeLabels = exclusiveLabel ex
. nodeLabels
. rootLabel
$ n
}
, subForest = map (filterExclusiveTree ex) . subForest $ n
}
exclusiveLabel :: Exclusivity -> Labels -> Labels
exclusiveLabel _ (Seq.null -> True) = Seq.empty
exclusiveLabel AllExclusive xs =
Seq.fromList . Set.toList . Set.fromList . F.toList $ xs
exclusiveLabel Exclusive xs =
if Seq.length uniqueSeq > 1
then Seq.empty
else uniqueSeq
where
uniqueSeq = Seq.fromList . Set.toList . Set.fromList . F.toList $ xs
exclusiveLabel Majority xs = Seq.singleton
. fst
. F.maximumBy (compare `on` snd)
. Map.toList
. Map.fromListWith (+)
. flip zip [1,1..]
. F.toList $ xs
addUniqueNodeIDs :: Tree NodeLabel -> Tree NodeLabel
addUniqueNodeIDs tree = fst . runState (go tree) $ 0
where
go :: Tree NodeLabel -> State Int (Tree NodeLabel)
go tree = do
newRootLabel <- updateRootLabel . rootLabel $ tree
newSubForest <- mapM go . subForest $ tree
return $ tree { rootLabel = newRootLabel
, subForest = newSubForest
}
updateRootLabel :: NodeLabel -> State Int NodeLabel
updateRootLabel n = do
nId <- get
modify (+ 1)
return $ n { nodeID = showt nId }