module Math.RandomTree.Label where
import Data.List
import Data.Maybe
import Data.Tree
import qualified Data.Map as M
import qualified Data.Sequence as Seq
import qualified Data.Foldable as F
import Control.Monad.Random
import System.Random.Shuffle
import Math.TreeFun.Types
modifyLabel :: (Eq a) => a -> a -> a -> a
modifyLabel old new n = if n == old
then new
else n
getNeighbors :: (Ord a) => Int -> a -> Tree (SuperNode a) -> [a]
getNeighbors neighborDistance l ( Node { rootLabel = SuperRoot
, subForest = ts } ) =
getNeighbors neighborDistance l
. head
. filter (M.member l . myLeaves . rootLabel)
$ ts
getNeighbors neighborDistance l ( Node { rootLabel = SuperNode { myRootLabel = _
, myParent = p
, myLeaves = ls }
, subForest = ts } )
| M.size ls == neighborDistance && relevant =
map fst . M.toAscList $ ls
| M.size ls > neighborDistance && relevant =
getNeighbors neighborDistance l
. head
. filter (M.member l . myLeaves . rootLabel)
$ ts
| M.size ls < neighborDistance && relevant && p == SuperRoot =
take neighborDistance
. (:) l
. filter (/= l)
. map fst
. M.toAscList
$ ls
| M.size ls < neighborDistance && relevant =
take neighborDistance
. (:) l
. filter (/= l)
. M.keys
. myLeaves
$ p
| otherwise = []
where
relevant = M.member l ls
clumpIt :: (Ord a, Eq b)
=> Int
-> Tree (SuperNode a)
-> a
-> b
-> MaybePropertyMap a b
-> MaybePropertyMap a b
clumpIt neighborDistance tree pointer property propertyMap =
F.foldl' (\acc x -> updateMap x property acc) propertyMap
$ neighbors pointer
where
updateMap k p = M.adjust
(\x -> if isNothing x then Just (Seq.singleton p) else x)
k
neighbors x = getNeighbors neighborDistance x tree
assignRandomClumpedProperties :: (Ord a, Eq b)
=> [b]
-> Int
-> Tree (SuperNode a)
-> StdGen
-> MaybePropertyMap a b
-> MaybePropertyMap a b
assignRandomClumpedProperties propertyList neighborDistance tree g propertyMap =
foldl' ( \acc (x, y)
-> clumpIt neighborDistance tree x y acc)
propertyMap
. zip shuffledLeaves
$ propertyList
where
shuffledLeaves = shuffle' (M.keys propertyMap) (M.size propertyMap) g
assignRandomProperties :: (Ord a)
=> [b]
-> StdGen
-> MaybePropertyMap a b
-> MaybePropertyMap a b
assignRandomProperties propertyList g propertyMap = M.fromList
. zip shuffledLeaves
. map (Just . Seq.singleton)
$ propertyList
where
shuffledLeaves = shuffle' (M.keys propertyMap) (M.size propertyMap) g
emptyPropertyMap :: (Ord a) => [a] -> MaybePropertyMap a b
emptyPropertyMap x = M.fromList . zip x . repeat $ Nothing
getPropertyMap :: (Ord a) => [a] -> PropertyMap a a
getPropertyMap x = M.fromList . zip x . map Seq.singleton $ x