module ELynx.Simulate.Coalescent
( simulate
) where
import Control.Monad.Primitive
import Data.Tree
import Statistics.Distribution
import System.Random.MWC
import ELynx.Data.Tree.MeasurableTree
import ELynx.Data.Tree.PhyloTree
import ELynx.Data.Tree.Tree
import ELynx.Distribution.CoalescentContinuous
simulate :: (PrimMonad m)
=> Int
-> Gen (PrimState m)
-> m (Tree PhyloIntLabel)
simulate n = simulate' n 0 trs
where trs = [ singleton (PhyloLabel i Nothing 0.0) | i <- [0..n-1] ]
simulate' :: (PrimMonad m)
=> Int
-> Int
-> [Tree PhyloIntLabel]
-> Gen (PrimState m)
-> m (Tree PhyloIntLabel)
simulate' n a trs g
| n <= 0 = error "Cannot construct trees without leaves."
| n == 1 && length trs /= 1 = error "Too many trees provided."
| n == 1 && length trs == 1 = return $ head trs
| otherwise =
do
i <- uniformR (1, n-1) g
t <- genContVar (coalescentDistributionCont n) g
let trs' = map (lengthenRoot t) trs
tl = trs' !! (i-1)
tr = trs' !! i
tm = Node (PhyloLabel a Nothing 0.0) [tl, tr]
trs'' = take (i-1) trs' ++ [tm] ++ drop (i+1) trs'
simulate' (n-1) a trs'' g