-- |
-- Module      :  ELynx.Tree.Simulate.Coalescent
-- Description :  Generate coalescent trees
-- Copyright   :  2021 Dominik Schrempf
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Wed May 16 13:13:11 2018.
module ELynx.Tree.Simulate.Coalescent
  ( simulate,
  )
where

import ELynx.Tree.Distribution.CoalescentContinuous
import ELynx.Tree.Length
import ELynx.Tree.Rooted
import Statistics.Distribution
import System.Random.Stateful

-- | Simulate a coalescent tree with @n@ leaves. The branch lengths are in units
-- of effective population size.
simulate ::
  StatefulGen g m =>
  -- | Number of leaves.
  Int ->
  g ->
  m (Tree Length Int)
simulate :: forall g (m :: * -> *).
StatefulGen g m =>
Int -> g -> m (Tree Length Int)
simulate Int
n = forall g (m :: * -> *).
StatefulGen g m =>
Int -> Int -> Forest Length Int -> g -> m (Tree Length Int)
simulate' Int
n Int
0 Forest Length Int
trs
  where
    trs :: Forest Length Int
trs = [forall e a. e -> a -> Forest e a -> Tree e a
Node Length
0 Int
i [] | Int
i <- [Int
0 .. Int
n forall a. Num a => a -> a -> a
- Int
1]]

simulate' ::
  StatefulGen g m =>
  Int ->
  Int ->
  Forest Length Int ->
  g ->
  m (Tree Length Int)
simulate' :: forall g (m :: * -> *).
StatefulGen g m =>
Int -> Int -> Forest Length Int -> g -> m (Tree Length Int)
simulate' Int
n Int
a Forest Length Int
trs g
g
  | Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 = forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot construct trees without leaves."
  | Int
n forall a. Eq a => a -> a -> Bool
== Int
1 Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => t a -> Int
length Forest Length Int
trs forall a. Eq a => a -> a -> Bool
/= Int
1 = forall a. HasCallStack => [Char] -> a
error [Char]
"Too many trees provided."
  | Int
n forall a. Eq a => a -> a -> Bool
== Int
1 Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => t a -> Int
length Forest Length Int
trs forall a. Eq a => a -> a -> Bool
== Int
1 = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
head Forest Length Int
trs
  | Bool
otherwise = do
      -- Indices of the leaves to join will be i-1 and i.
      Int
i <- forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
uniformRM (Int
1, Int
n forall a. Num a => a -> a -> a
- Int
1) g
g
      -- The time of the coalescent event.
      Length
t <- Double -> Length
toLengthUnsafe forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall d g (m :: * -> *).
(ContGen d, StatefulGen g m) =>
d -> g -> m Double
genContVar (Int -> ExponentialDistribution
coalescentDistributionCont Int
n) g
g
      let trs' :: Forest Length Int
trs' = forall a b. (a -> b) -> [a] -> [b]
map (forall e a. (e -> e) -> Tree e a -> Tree e a
modifyStem (forall a. Num a => a -> a -> a
+ Length
t)) Forest Length Int
trs -- Move time 't' up on the tree.
          tl :: Tree Length Int
tl = Forest Length Int
trs' forall a. [a] -> Int -> a
!! (Int
i forall a. Num a => a -> a -> a
- Int
1)
          tr :: Tree Length Int
tr = Forest Length Int
trs' forall a. [a] -> Int -> a
!! Int
i
          -- Join the two chosen trees.
          tm :: Tree Length Int
tm = forall e a. e -> a -> Forest e a -> Tree e a
Node Length
0 Int
a [Tree Length Int
tl, Tree Length Int
tr]
          -- Take the trees on the left, the merged tree, and the trees on the right.
          trs'' :: Forest Length Int
trs'' = forall a. Int -> [a] -> [a]
take (Int
i forall a. Num a => a -> a -> a
- Int
1) Forest Length Int
trs' forall a. [a] -> [a] -> [a]
++ [Tree Length Int
tm] forall a. [a] -> [a] -> [a]
++ forall a. Int -> [a] -> [a]
drop (Int
i forall a. Num a => a -> a -> a
+ Int
1) Forest Length Int
trs'
      forall g (m :: * -> *).
StatefulGen g m =>
Int -> Int -> Forest Length Int -> g -> m (Tree Length Int)
simulate' (Int
n forall a. Num a => a -> a -> a
- Int
1) Int
a Forest Length Int
trs'' g
g