{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}

module Epidemic.Utility where

import Control.Applicative
import Control.Monad (liftM)
import Control.Monad.Primitive (PrimMonad, PrimState)
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as Char8
import qualified Data.List as List
import qualified Data.Maybe as Maybe
import qualified Data.Vector as V
import Epidemic
import Epidemic.Types.Events
import Epidemic.Types.Parameter
import Epidemic.Types.Population
import Epidemic.Types.Simulation
import Epidemic.Types.Time
  ( AbsoluteTime(..)
  , Timed(..)
  , TimeDelta(..)
  , diracDeltaValue
  , nextTime
  , cadlagValue
  , timeAfterDelta
  )
import GHC.Generics (Generic)
import System.Random.MWC
import System.Random.MWC.Distributions (exponential)


initialIdentifier :: Identifier
initialIdentifier :: Identifier
initialIdentifier = Integer -> Identifier
Identifier Integer
1

newPerson :: Identifier -> (Person, Identifier)
newPerson :: Identifier -> (Person, Identifier)
newPerson idntty :: Identifier
idntty@(Identifier Integer
idInt) = (Identifier -> Person
Person Identifier
idntty, Integer -> Identifier
Identifier (Integer
idInt Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1))

selectElem :: V.Vector a -> Int -> (a, V.Vector a)
selectElem :: Vector a -> Int -> (a, Vector a)
selectElem Vector a
v Int
n
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = (Vector a -> a
forall a. Vector a -> a
V.head Vector a
v, Vector a -> Vector a
forall a. Vector a -> Vector a
V.tail Vector a
v)
  | Bool
otherwise =
    let (Vector a
foo, Vector a
bar) = Int -> Vector a -> (Vector a, Vector a)
forall a. Int -> Vector a -> (Vector a, Vector a)
V.splitAt Int
n Vector a
v
     in (Vector a -> a
forall a. Vector a -> a
V.head Vector a
bar, Vector a
foo Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
V.++ (Vector a -> Vector a
forall a. Vector a -> Vector a
V.tail Vector a
bar))

randomPerson :: People -> GenIO -> IO (Person, People)
randomPerson :: People -> GenIO -> IO (Person, People)
randomPerson people :: People
people@(People Vector Person
persons) GenIO
gen = do
  Double
u <- GenIO -> IO Double
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
Gen (PrimState m) -> m a
uniform GenIO
gen
  let personIx :: Int
personIx = Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double
u Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Double) -> Int -> Double
forall a b. (a -> b) -> a -> b
$ People -> Int
numPeople People
people :: Double))
      (Person
person, Vector Person
remPeople) = Vector Person -> Int -> (Person, Vector Person)
forall a. Vector a -> Int -> (a, Vector a)
selectElem Vector Person
persons Int
personIx
   in (Person, People) -> IO (Person, People)
forall (m :: * -> *) a. Monad m => a -> m a
return (Person
person, Vector Person -> People
People Vector Person
remPeople)

type NName = Maybe String

type NLength = Maybe Double

data NBranch =
  NBranch NSubtree NLength
  deriving (NBranch -> NBranch -> Bool
(NBranch -> NBranch -> Bool)
-> (NBranch -> NBranch -> Bool) -> Eq NBranch
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NBranch -> NBranch -> Bool
$c/= :: NBranch -> NBranch -> Bool
== :: NBranch -> NBranch -> Bool
$c== :: NBranch -> NBranch -> Bool
Eq)

instance Show NBranch where
  show :: NBranch -> String
show (NBranch NSubtree
st (Just Double
l)) = NSubtree -> String
forall a. Show a => a -> String
show NSubtree
st String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
":" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Double -> String
forall a. Show a => a -> String
show Double
l
  show (NBranch NSubtree
st Maybe Double
Nothing) = NSubtree -> String
forall a. Show a => a -> String
show NSubtree
st

data NBranchSet =
  NBranchSet [NBranch]
  deriving (NBranchSet -> NBranchSet -> Bool
(NBranchSet -> NBranchSet -> Bool)
-> (NBranchSet -> NBranchSet -> Bool) -> Eq NBranchSet
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NBranchSet -> NBranchSet -> Bool
$c/= :: NBranchSet -> NBranchSet -> Bool
== :: NBranchSet -> NBranchSet -> Bool
$c== :: NBranchSet -> NBranchSet -> Bool
Eq)

instance Show NBranchSet where
  show :: NBranchSet -> String
show (NBranchSet [NBranch]
bs) = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ (String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
List.intercalate String
"," ((NBranch -> String) -> [NBranch] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map NBranch -> String
forall a. Show a => a -> String
show [NBranch]
bs)) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

data NSubtree
  = NLeaf NName
  | NInternal NBranchSet
  deriving (NSubtree -> NSubtree -> Bool
(NSubtree -> NSubtree -> Bool)
-> (NSubtree -> NSubtree -> Bool) -> Eq NSubtree
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NSubtree -> NSubtree -> Bool
$c/= :: NSubtree -> NSubtree -> Bool
== :: NSubtree -> NSubtree -> Bool
$c== :: NSubtree -> NSubtree -> Bool
Eq)

instance Show NSubtree where
  show :: NSubtree -> String
show (NLeaf (Just String
n)) = String
n
  show (NLeaf Maybe String
Nothing) = String
""
  show (NInternal NBranchSet
bs) = NBranchSet -> String
forall a. Show a => a -> String
show NBranchSet
bs

data NTree =
  NTree [NBranch]
  deriving (NTree -> NTree -> Bool
(NTree -> NTree -> Bool) -> (NTree -> NTree -> Bool) -> Eq NTree
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: NTree -> NTree -> Bool
$c/= :: NTree -> NTree -> Bool
== :: NTree -> NTree -> Bool
$c== :: NTree -> NTree -> Bool
Eq)

instance Show NTree where
  show :: NTree -> String
show (NTree [NBranch]
bs) = NBranchSet -> String
forall a. Show a => a -> String
show ([NBranch] -> NBranchSet
NBranchSet [NBranch]
bs) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
";"

-- | Example run
--   > (Success foo) = parseString newickTree mempty "((foo:1.1,bar:1.2):1.3,baz:1.4);"
--   > (Success bar) = parseString newickTree mempty $ show foo
--   > foo == bar
--   True
sort :: Ord a => [a] -> [a]
sort :: [a] -> [a]
sort = [a] -> [a]
forall a. Ord a => [a] -> [a]
List.sort

count' :: (a -> Bool) -> [a] -> Int
count' :: (a -> Bool) -> [a] -> Int
count' a -> Bool
p = Int -> [a] -> Int
forall t. Num t => t -> [a] -> t
go Int
0
  where
    go :: t -> [a] -> t
go t
n [] = t
n
    go t
n (a
x:[a]
xs)
      | a -> Bool
p a
x = t -> [a] -> t
go (t
n t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) [a]
xs
      | Bool
otherwise = t -> [a] -> t
go t
n [a]
xs

-- | Run a simulation described by a configuration object with the provided
-- PRNG.
simulationWithGenIO ::
     (ModelParameters a b, Population b)
  => SimulationConfiguration a b
  -> (a -> AbsoluteTime -> Maybe (b -> Bool) -> SimulationState b -> GenIO -> IO (SimulationState b))
  -> GenIO
  -> IO [EpidemicEvent]
simulationWithGenIO :: SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> GenIO
-> IO [EpidemicEvent]
simulationWithGenIO config :: SimulationConfiguration a b
config@SimulationConfiguration {a
b
Bool
Maybe (b -> Bool)
Identifier
TimeDelta
AbsoluteTime
scRequireCherry :: forall r p. SimulationConfiguration r p -> Bool
scValidPopulation :: forall r p. SimulationConfiguration r p -> Maybe (p -> Bool)
scSimDuration :: forall r p. SimulationConfiguration r p -> TimeDelta
scStartTime :: forall r p. SimulationConfiguration r p -> AbsoluteTime
scNewIdentifier :: forall r p. SimulationConfiguration r p -> Identifier
scPopulation :: forall r p. SimulationConfiguration r p -> p
scRates :: forall r p. SimulationConfiguration r p -> r
scRequireCherry :: Bool
scValidPopulation :: Maybe (b -> Bool)
scSimDuration :: TimeDelta
scStartTime :: AbsoluteTime
scNewIdentifier :: Identifier
scPopulation :: b
scRates :: a
..} a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEventsFunc GenIO
gen =
  if Bool
scRequireCherry
    then do
      SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> GenIO
-> IO [EpidemicEvent]
forall a b.
(ModelParameters a b, Population b) =>
SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> GenIO
-> IO [EpidemicEvent]
simulation' SimulationConfiguration a b
config a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEventsFunc GenIO
gen
    else do
      SimulationState (AbsoluteTime
_, [EpidemicEvent]
events, b
_, Identifier
_) <-
        a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEventsFunc
          a
scRates
          (AbsoluteTime -> TimeDelta -> AbsoluteTime
timeAfterDelta AbsoluteTime
scStartTime TimeDelta
scSimDuration)
          Maybe (b -> Bool)
scValidPopulation
          ((AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
forall b.
(AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
SimulationState (Double -> AbsoluteTime
AbsoluteTime Double
0, [], b
scPopulation, Identifier
scNewIdentifier))
          GenIO
gen
      [EpidemicEvent] -> IO [EpidemicEvent]
forall (m :: * -> *) a. Monad m => a -> m a
return ([EpidemicEvent] -> IO [EpidemicEvent])
-> [EpidemicEvent] -> IO [EpidemicEvent]
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> [EpidemicEvent]
forall a. Ord a => [a] -> [a]
sort [EpidemicEvent]
events

-- | Run a simulation described by a configuration object using the fixed PRNG
-- that is hardcoded in the @mwc-random@ package.
simulation ::
     (ModelParameters a b, Population b)
  => SimulationConfiguration a b
  -> (a -> AbsoluteTime -> Maybe (b -> Bool) -> SimulationState b -> GenIO -> IO (SimulationState b))
  -> IO [EpidemicEvent]
simulation :: SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> IO [EpidemicEvent]
simulation SimulationConfiguration a b
config a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEventsFunc = do
  Gen RealWorld
gen <- IO GenIO
forall (m :: * -> *). PrimMonad m => m (Gen (PrimState m))
System.Random.MWC.create :: IO GenIO
  SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> GenIO
-> IO [EpidemicEvent]
forall a b.
(ModelParameters a b, Population b) =>
SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> GenIO
-> IO [EpidemicEvent]
simulationWithGenIO SimulationConfiguration a b
config a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEventsFunc Gen RealWorld
GenIO
gen

-- | Predicate for whether an epidemic event will appear as a leaf in the
-- reconstructed tree.
isReconTreeLeaf :: EpidemicEvent -> Bool
isReconTreeLeaf :: EpidemicEvent -> Bool
isReconTreeLeaf EpidemicEvent
e =
  case EpidemicEvent
e of
    IndividualSample {Bool
Person
AbsoluteTime
indSampSeq :: EpidemicEvent -> Bool
indSampPerson :: EpidemicEvent -> Person
indSampTime :: EpidemicEvent -> AbsoluteTime
indSampSeq :: Bool
indSampPerson :: Person
indSampTime :: AbsoluteTime
..} -> Bool
indSampSeq
    PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
..} -> Bool
popSampSeq
    EpidemicEvent
_ -> Bool
False

-- | Simulation conditioned upon there being at least two sequenced samples.
-- NOTE This function is deprecated and will be removed in future versions.
simulation' ::
     (ModelParameters a b, Population b)
  => SimulationConfiguration a b
  -> (a -> AbsoluteTime -> Maybe (b -> Bool) -> SimulationState b -> GenIO -> IO (SimulationState b))
  -> GenIO
  -> IO [EpidemicEvent]
simulation' :: SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> GenIO
-> IO [EpidemicEvent]
simulation' config :: SimulationConfiguration a b
config@SimulationConfiguration {a
b
Bool
Maybe (b -> Bool)
Identifier
TimeDelta
AbsoluteTime
scRequireCherry :: Bool
scValidPopulation :: Maybe (b -> Bool)
scSimDuration :: TimeDelta
scStartTime :: AbsoluteTime
scNewIdentifier :: Identifier
scPopulation :: b
scRates :: a
scRequireCherry :: forall r p. SimulationConfiguration r p -> Bool
scValidPopulation :: forall r p. SimulationConfiguration r p -> Maybe (p -> Bool)
scSimDuration :: forall r p. SimulationConfiguration r p -> TimeDelta
scStartTime :: forall r p. SimulationConfiguration r p -> AbsoluteTime
scNewIdentifier :: forall r p. SimulationConfiguration r p -> Identifier
scPopulation :: forall r p. SimulationConfiguration r p -> p
scRates :: forall r p. SimulationConfiguration r p -> r
..} a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEventsFunc GenIO
gen = do
  SimulationState (AbsoluteTime
_, [EpidemicEvent]
events, b
_, Identifier
_) <-
    a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEventsFunc
      a
scRates
      (AbsoluteTime -> TimeDelta -> AbsoluteTime
timeAfterDelta AbsoluteTime
scStartTime TimeDelta
scSimDuration)
      Maybe (b -> Bool)
scValidPopulation
      ((AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
forall b.
(AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
SimulationState (Double -> AbsoluteTime
AbsoluteTime Double
0, [], b
scPopulation, Identifier
scNewIdentifier))
      GenIO
gen
  if (EpidemicEvent -> Bool) -> [EpidemicEvent] -> Int
forall a. (a -> Bool) -> [a] -> Int
count' EpidemicEvent -> Bool
isReconTreeLeaf [EpidemicEvent]
events Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2
    then [EpidemicEvent] -> IO [EpidemicEvent]
forall (m :: * -> *) a. Monad m => a -> m a
return ([EpidemicEvent] -> IO [EpidemicEvent])
-> [EpidemicEvent] -> IO [EpidemicEvent]
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> [EpidemicEvent]
forall a. Ord a => [a] -> [a]
sort [EpidemicEvent]
events
    else SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> GenIO
-> IO [EpidemicEvent]
forall a b.
(ModelParameters a b, Population b) =>
SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> GenIO
-> IO [EpidemicEvent]
simulation' SimulationConfiguration a b
config a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEventsFunc GenIO
gen

-- | Run a simulation described by a configuration object but using a random
-- seed generated by the system rather than a seed
simulationWithSystemRandom ::
     (ModelParameters a b, Population b)
  => SimulationConfiguration a b
  -> (a -> AbsoluteTime -> Maybe (b -> Bool) -> SimulationState b -> GenIO -> IO (SimulationState b))
  -> IO [EpidemicEvent]
simulationWithSystemRandom :: SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> IO [EpidemicEvent]
simulationWithSystemRandom config :: SimulationConfiguration a b
config@SimulationConfiguration {a
b
Bool
Maybe (b -> Bool)
Identifier
TimeDelta
AbsoluteTime
scRequireCherry :: Bool
scValidPopulation :: Maybe (b -> Bool)
scSimDuration :: TimeDelta
scStartTime :: AbsoluteTime
scNewIdentifier :: Identifier
scPopulation :: b
scRates :: a
scRequireCherry :: forall r p. SimulationConfiguration r p -> Bool
scValidPopulation :: forall r p. SimulationConfiguration r p -> Maybe (p -> Bool)
scSimDuration :: forall r p. SimulationConfiguration r p -> TimeDelta
scStartTime :: forall r p. SimulationConfiguration r p -> AbsoluteTime
scNewIdentifier :: forall r p. SimulationConfiguration r p -> Identifier
scPopulation :: forall r p. SimulationConfiguration r p -> p
scRates :: forall r p. SimulationConfiguration r p -> r
..} a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEventsFunc = do
  SimulationState (AbsoluteTime
_, [EpidemicEvent]
events, b
_, Identifier
_) <-
    (GenIO -> IO (SimulationState b)) -> IO (SimulationState b)
forall (m :: * -> *) a.
PrimBase m =>
(Gen (PrimState m) -> m a) -> IO a
withSystemRandom ((GenIO -> IO (SimulationState b)) -> IO (SimulationState b))
-> (GenIO -> IO (SimulationState b)) -> IO (SimulationState b)
forall a b. (a -> b) -> a -> b
$ \GenIO
g ->
      a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEventsFunc
        a
scRates
        (AbsoluteTime -> TimeDelta -> AbsoluteTime
timeAfterDelta AbsoluteTime
scStartTime TimeDelta
scSimDuration)
        Maybe (b -> Bool)
scValidPopulation
        ((AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
forall b.
(AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
SimulationState (Double -> AbsoluteTime
AbsoluteTime Double
0, [], b
scPopulation, Identifier
scNewIdentifier))
        GenIO
g
  if Bool
scRequireCherry
    then (if (EpidemicEvent -> Bool) -> [EpidemicEvent] -> Int
forall a. (a -> Bool) -> [a] -> Int
count' EpidemicEvent -> Bool
isReconTreeLeaf [EpidemicEvent]
events Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2
            then [EpidemicEvent] -> IO [EpidemicEvent]
forall (m :: * -> *) a. Monad m => a -> m a
return ([EpidemicEvent] -> IO [EpidemicEvent])
-> [EpidemicEvent] -> IO [EpidemicEvent]
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> [EpidemicEvent]
forall a. Ord a => [a] -> [a]
sort [EpidemicEvent]
events
            else SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> IO [EpidemicEvent]
forall a b.
(ModelParameters a b, Population b) =>
SimulationConfiguration a b
-> (a
    -> AbsoluteTime
    -> Maybe (b -> Bool)
    -> SimulationState b
    -> GenIO
    -> IO (SimulationState b))
-> IO [EpidemicEvent]
simulationWithSystemRandom SimulationConfiguration a b
config a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEventsFunc)
    else [EpidemicEvent] -> IO [EpidemicEvent]
forall (m :: * -> *) a. Monad m => a -> m a
return ([EpidemicEvent] -> IO [EpidemicEvent])
-> [EpidemicEvent] -> IO [EpidemicEvent]
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> [EpidemicEvent]
forall a. Ord a => [a] -> [a]
sort [EpidemicEvent]
events

-- | The number of lineages at the end of a simulation.
finalSize ::
     [EpidemicEvent] -- ^ The events from the simulation
  -> Integer
finalSize :: [EpidemicEvent] -> Integer
finalSize = (Integer -> EpidemicEvent -> Integer)
-> Integer -> [EpidemicEvent] -> Integer
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Integer
x EpidemicEvent
y -> Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ EpidemicEvent -> Integer
eventPopDelta EpidemicEvent
y) Integer
1

-- | Generate exponentially distributed random variates with inhomogeneous rate
-- starting from a particular point in time.
--
-- Assuming the @stepFunc@ is the intensity of arrivals and @t0@ is the start
-- time this returns @t1@ the time of the next arrival.
inhomExponential ::
     PrimMonad m
  => Timed Double -- ^ Step function
  -> AbsoluteTime -- ^ Start time
  -> Gen (PrimState m) -- ^ Generator
  -> m (Maybe AbsoluteTime)
inhomExponential :: Timed Double
-> AbsoluteTime -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
inhomExponential Timed Double
stepFunc AbsoluteTime
t0 = AbsoluteTime
-> Timed Double -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
forall (m :: * -> *).
PrimMonad m =>
AbsoluteTime
-> Timed Double -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
randInhomExp AbsoluteTime
t0 Timed Double
stepFunc

-- | Generate exponentially distributed random variates with inhomogeneous rate.
--
-- __TODO__ The algorithm used here generates more variates than are needed. It
-- would be nice to use a more efficient implementation.
--
randInhomExp ::
     PrimMonad m
  => AbsoluteTime -- ^ Timer
  -> Timed Double -- ^ Step function
  -> Gen (PrimState m) -- ^ Generator.
  -> m (Maybe AbsoluteTime)
randInhomExp :: AbsoluteTime
-> Timed Double -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
randInhomExp AbsoluteTime
crrT Timed Double
stepFunc Gen (PrimState m)
gen =
  let crrR :: Maybe Double
crrR = Timed Double -> AbsoluteTime -> Maybe Double
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Double
stepFunc AbsoluteTime
crrT
      nxtT :: Maybe AbsoluteTime
nxtT = Timed Double -> AbsoluteTime -> Maybe AbsoluteTime
forall a. Timed a -> AbsoluteTime -> Maybe AbsoluteTime
nextTime Timed Double
stepFunc AbsoluteTime
crrT
   in if Maybe Double -> Bool
forall a. Maybe a -> Bool
Maybe.isJust Maybe Double
crrR Bool -> Bool -> Bool
&& Maybe AbsoluteTime -> Bool
forall a. Maybe a -> Bool
Maybe.isJust Maybe AbsoluteTime
nxtT
        then do
          Double
crrD <- Double -> Gen (PrimState m) -> m Double
forall g (m :: * -> *). StatefulGen g m => Double -> g -> m Double
exponential (Maybe Double -> Double
forall a. HasCallStack => Maybe a -> a
Maybe.fromJust Maybe Double
crrR) Gen (PrimState m)
gen
          let propT :: AbsoluteTime
propT = AbsoluteTime -> TimeDelta -> AbsoluteTime
timeAfterDelta AbsoluteTime
crrT (Double -> TimeDelta
TimeDelta Double
crrD)
          if AbsoluteTime
propT AbsoluteTime -> AbsoluteTime -> Bool
forall a. Ord a => a -> a -> Bool
< Maybe AbsoluteTime -> AbsoluteTime
forall a. HasCallStack => Maybe a -> a
Maybe.fromJust Maybe AbsoluteTime
nxtT
            then Maybe AbsoluteTime -> m (Maybe AbsoluteTime)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe AbsoluteTime -> m (Maybe AbsoluteTime))
-> Maybe AbsoluteTime -> m (Maybe AbsoluteTime)
forall a b. (a -> b) -> a -> b
$ AbsoluteTime -> Maybe AbsoluteTime
forall a. a -> Maybe a
Just AbsoluteTime
propT
            else AbsoluteTime
-> Timed Double -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
forall (m :: * -> *).
PrimMonad m =>
AbsoluteTime
-> Timed Double -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
randInhomExp (Maybe AbsoluteTime -> AbsoluteTime
forall a. HasCallStack => Maybe a -> a
Maybe.fromJust Maybe AbsoluteTime
nxtT) Timed Double
stepFunc Gen (PrimState m)
gen
        else Maybe AbsoluteTime -> m (Maybe AbsoluteTime)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe AbsoluteTime
forall a. Maybe a
Nothing

-- | Helper function for converting between the Maybe monad and the Either
-- monad.
maybeToRight :: a -> Maybe b -> Either a b
maybeToRight :: a -> Maybe b -> Either a b
maybeToRight a
a Maybe b
maybeB =
  case Maybe b
maybeB of
    (Just b
b) -> b -> Either a b
forall a b. b -> Either a b
Right b
b
    Maybe b
Nothing -> a -> Either a b
forall a b. a -> Either a b
Left a
a