module TreeTransform
( convertToSuperTree
, getPropertyMap
, innerToLeaves
, filterExclusiveTree
) where
import qualified Data.Map.Strict as Map
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import Data.Tree
import qualified Data.Foldable as F
import Data.Function (on)
import qualified Data.Text as T
import Math.TreeFun.Types
import Math.TreeFun.Tree
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