{-# LANGUAGE RecordWildCards #-}

-- |
-- Module: Epidemic
-- Copyright: (c) 2021 Alexander E. Zarebski
-- License: MIT
--
-- Maintainer: Alexander E. Zarebski <aezarebski@gmail.com>
-- Stability: unstable
-- Portability: ghc
--
-- This package provides functionality for simulating stochastic epidemic
-- models, in particular those that are of interest in phylodynamics. There are
-- several models provided by the package, eg @Epidemic.Model.BDSCOD@, however
-- there should be the basic functionality to implement a wide range of models
-- available. Each of the models included in the package provide a
-- @configuration@ function which can be used to get a 'SimulationConfiguration'
-- and a @randomEvent@ function which returns a 'SimulationRandEvent'. With
-- these you can then use 'allEvents' to get all of the events in a simulation.
--
-- This package also provides some functionality for working with observation
-- models, both epidemiological and phylogenetic. 'Observation' values are used
-- to describe the possible observation of an 'EpidemicEvent'.
--
-- There is an example of how to use this package in the documentation for
-- "Epidemic.Model.InhomogeneousBDSCOD".

module Epidemic where

import           Data.List                 (nub)
import           Data.Maybe                (fromJust, isJust, isNothing)
import qualified Data.Vector               as V
import           Epidemic.Types.Events
import           Epidemic.Types.Parameter
import           Epidemic.Types.Population
import           Epidemic.Types.Simulation (SimulationRandEvent (..),
                                            SimulationState (..),
                                            TerminationHandler (..))
import           Epidemic.Types.Time       (AbsoluteTime (..), Timed (..),
                                            diracDeltaValue, nextTime)
import           System.Random.MWC

-- | The number of people added or removed in an event. In the case of an
-- extinction event the number of people removed is arbitrarily set to zero
-- because this information is available from the prior event in the sequence.
eventPopDelta :: EpidemicEvent -> Integer
eventPopDelta :: EpidemicEvent -> Integer
eventPopDelta EpidemicEvent
e =
  case EpidemicEvent
e of
    Infection {}          -> Integer
1
    Removal {}            -> -Integer
1
    IndividualSample {}   -> -Integer
1
    PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
..} -> Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ People -> Int
numPeople People
popSampPeople
    StoppingTime {}       -> Integer
0
    Extinction {}         -> Integer
0

-- | The first scheduled event after a given time.
firstScheduled ::
     AbsoluteTime -- ^ The given time
  -> Timed Probability -- ^ The information about all scheduled events
  -> Maybe (AbsoluteTime, Probability)
firstScheduled :: AbsoluteTime
-> Timed Probability -> Maybe (AbsoluteTime, Probability)
firstScheduled AbsoluteTime
time Timed Probability
timedProb = do
  AbsoluteTime
time' <- Timed Probability -> AbsoluteTime -> Maybe AbsoluteTime
forall a. Timed a -> AbsoluteTime -> Maybe AbsoluteTime
nextTime Timed Probability
timedProb AbsoluteTime
time
  Probability
prob' <- Timed Probability -> AbsoluteTime -> Maybe Probability
forall a. Timed a -> AbsoluteTime -> Maybe a
diracDeltaValue Timed Probability
timedProb AbsoluteTime
time'
  (AbsoluteTime, Probability) -> Maybe (AbsoluteTime, Probability)
forall (m :: * -> *) a. Monad m => a -> m a
return (AbsoluteTime
time', Probability
prob')

-- | Predicate for whether there is a scheduled event during an interval. NOTE
-- that this does not consider events that happen at the start of the interval
-- as occurring between the times.
--
-- >>> tA = AbsoluteTime 1.0
-- >>> tB = AbsoluteTime 2.0
-- >>> noScheduledEvent tA tB <$> asTimed [(AbsoluteTime 1.5, 0.5)]
-- Just False
-- >>> noScheduledEvent tA tB <$> asTimed [(AbsoluteTime 2.5, 0.5)]
-- Just True
-- >>> noScheduledEvent tA tB <$> asTimed [(tA, 0.5)]
-- Just True
--
noScheduledEvent ::
     AbsoluteTime -- ^ Start time for interval
  -> AbsoluteTime -- ^ End time for interval
  -> Timed Probability -- ^ Information about all scheduled events
  -> Bool
noScheduledEvent :: AbsoluteTime -> AbsoluteTime -> Timed Probability -> Bool
noScheduledEvent AbsoluteTime
_ AbsoluteTime
_ (Timed []) = Bool
True
noScheduledEvent AbsoluteTime
a AbsoluteTime
b (Timed ((AbsoluteTime
shedTime, Probability
_):[(AbsoluteTime, Probability)]
scheduledEvents)) =
  Bool -> Bool
not (AbsoluteTime
a AbsoluteTime -> AbsoluteTime -> Bool
forall a. Ord a => a -> a -> Bool
< AbsoluteTime
shedTime Bool -> Bool -> Bool
&& AbsoluteTime
shedTime AbsoluteTime -> AbsoluteTime -> Bool
forall a. Ord a => a -> a -> Bool
<= AbsoluteTime
b) Bool -> Bool -> Bool
&&
  AbsoluteTime -> AbsoluteTime -> Timed Probability -> Bool
noScheduledEvent AbsoluteTime
a AbsoluteTime
b ([(AbsoluteTime, Probability)] -> Timed Probability
forall a. [(AbsoluteTime, a)] -> Timed a
Timed [(AbsoluteTime, Probability)]
scheduledEvents)

-- | A list of the people involved in an 'EpidemicEvent'.
personsInEvent :: EpidemicEvent -> [Person]
personsInEvent :: EpidemicEvent -> [Person]
personsInEvent EpidemicEvent
e =
  case EpidemicEvent
e of
    Infection AbsoluteTime
_ Person
p1 Person
p2 -> [Person
p1, Person
p2]
    Removal AbsoluteTime
_ Person
p -> [Person
p]
    IndividualSample {Bool
Person
AbsoluteTime
indSampSeq :: EpidemicEvent -> Bool
indSampPerson :: EpidemicEvent -> Person
indSampTime :: EpidemicEvent -> AbsoluteTime
indSampSeq :: Bool
indSampPerson :: Person
indSampTime :: AbsoluteTime
..} -> [Person
indSampPerson]
    PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
..} -> Vector Person -> [Person]
forall a. Vector a -> [a]
V.toList Vector Person
personVec
      where (People Vector Person
personVec) = People
popSampPeople
    Extinction {} -> []
    StoppingTime {} -> []

peopleInEvents :: [EpidemicEvent] -> People
peopleInEvents :: [EpidemicEvent] -> People
peopleInEvents [EpidemicEvent]
events =
  Vector Person -> People
People (Vector Person -> People)
-> ([[Person]] -> Vector Person) -> [[Person]] -> People
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Person] -> Vector Person
forall a. [a] -> Vector a
V.fromList ([Person] -> Vector Person)
-> ([[Person]] -> [Person]) -> [[Person]] -> Vector Person
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Person] -> [Person]
forall a. Eq a => [a] -> [a]
nub ([Person] -> [Person])
-> ([[Person]] -> [Person]) -> [[Person]] -> [Person]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Person]] -> [Person]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Person]] -> People) -> [[Person]] -> People
forall a b. (a -> b) -> a -> b
$ (EpidemicEvent -> [Person]) -> [EpidemicEvent] -> [[Person]]
forall a b. (a -> b) -> [a] -> [b]
map EpidemicEvent -> [Person]
personsInEvent [EpidemicEvent]
events

-- | Predicate for whether the first person infected the second in the given event
infected ::
     Person -- ^ Potential infector
  -> Person -- ^ Potential infectee
  -> EpidemicEvent -- ^ Given event
  -> Bool
infected :: Person -> Person -> EpidemicEvent -> Bool
infected Person
p1 Person
p2 EpidemicEvent
e =
  case EpidemicEvent
e of
    (Infection AbsoluteTime
_ Person
infector Person
infectee) -> Person
infector Person -> Person -> Bool
forall a. Eq a => a -> a -> Bool
== Person
p1 Bool -> Bool -> Bool
&& Person
infectee Person -> Person -> Bool
forall a. Eq a => a -> a -> Bool
== Person
p2
    EpidemicEvent
_                               -> Bool
False

-- | The people infected by a particular person in a list of events.
infectedBy ::
     Person -- ^ Potential infector
  -> [EpidemicEvent] -- ^ Events
  -> People
infectedBy :: Person -> [EpidemicEvent] -> People
infectedBy Person
person [EpidemicEvent]
events =
  case [EpidemicEvent]
events of
    [] -> Vector Person -> People
People Vector Person
forall a. Vector a
V.empty
    (Infection AbsoluteTime
_ Person
infector Person
infectee:[EpidemicEvent]
es) ->
      if Person
infector Person -> Person -> Bool
forall a. Eq a => a -> a -> Bool
== Person
person
        then Person -> People -> People
addPerson Person
infectee (People -> People) -> People -> People
forall a b. (a -> b) -> a -> b
$ Person -> [EpidemicEvent] -> People
infectedBy Person
person [EpidemicEvent]
es
        else Person -> [EpidemicEvent] -> People
infectedBy Person
person [EpidemicEvent]
es
    (EpidemicEvent
_:[EpidemicEvent]
es) -> Person -> [EpidemicEvent] -> People
infectedBy Person
person [EpidemicEvent]
es

-- | Run the simulation until the specified stopping time and return a
-- @SimulationState@ which holds the history of the simulation.
allEvents ::
     (ModelParameters a b, Population b)
  => SimulationRandEvent a b
  -> a
  -> AbsoluteTime -- ^ time at which to stop the simulation
  -> Maybe (TerminationHandler b c)
  -> SimulationState b c -- ^ the initial/current state of the simulation
  -> GenIO
  -> IO (SimulationState b c)
allEvents :: SimulationRandEvent a b
-> a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEvents SimulationRandEvent a b
_ a
_ AbsoluteTime
_ Maybe (TerminationHandler b c)
_ ts :: SimulationState b c
ts@(TerminatedSimulation Maybe c
_) GenIO
_ = SimulationState b c -> IO (SimulationState b c)
forall (m :: * -> *) a. Monad m => a -> m a
return SimulationState b c
ts
allEvents (SimulationRandEvent a
-> AbsoluteTime
-> b
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, b, Identifier)
randEvent) a
modelParams AbsoluteTime
maxTime Maybe (TerminationHandler b c)
maybeTermHandler (SimulationState (AbsoluteTime
currTime, [EpidemicEvent]
currEvents, b
currPop, Identifier
currId)) GenIO
gen =
  let isNotTerminated :: b -> Bool
isNotTerminated = case Maybe (TerminationHandler b c)
maybeTermHandler of
        Maybe (TerminationHandler b c)
Nothing                                   -> Bool -> b -> Bool
forall a b. a -> b -> a
const Bool
True
        Just (TerminationHandler b -> Bool
hasTerminated [EpidemicEvent] -> c
_) -> Bool -> Bool
not (Bool -> Bool) -> (b -> Bool) -> b -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Bool
hasTerminated
  in if b -> Bool
isNotTerminated b
currPop
     then if b -> Bool
forall a. Population a => a -> Bool
isInfected b
currPop
           then do
             (AbsoluteTime
newTime, EpidemicEvent
event, b
newPop, Identifier
newId) <-
               a
-> AbsoluteTime
-> b
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, b, Identifier)
randEvent a
modelParams AbsoluteTime
currTime b
currPop Identifier
currId GenIO
gen
             if AbsoluteTime
newTime AbsoluteTime -> AbsoluteTime -> Bool
forall a. Ord a => a -> a -> Bool
< AbsoluteTime
maxTime
               then SimulationRandEvent a b
-> a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
forall a b c.
(ModelParameters a b, Population b) =>
SimulationRandEvent a b
-> a
-> AbsoluteTime
-> Maybe (TerminationHandler b c)
-> SimulationState b c
-> GenIO
-> IO (SimulationState b c)
allEvents
                      ((a
 -> AbsoluteTime
 -> b
 -> Identifier
 -> GenIO
 -> IO (AbsoluteTime, EpidemicEvent, b, Identifier))
-> SimulationRandEvent a b
forall a b.
(ModelParameters a b, Population b) =>
(a
 -> AbsoluteTime
 -> b
 -> Identifier
 -> GenIO
 -> IO (AbsoluteTime, EpidemicEvent, b, Identifier))
-> SimulationRandEvent a b
SimulationRandEvent a
-> AbsoluteTime
-> b
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, b, Identifier)
randEvent)
                      a
modelParams
                      AbsoluteTime
maxTime
                      Maybe (TerminationHandler b c)
maybeTermHandler
                      ((AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
forall b c.
(AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
SimulationState
                         (AbsoluteTime
newTime, EpidemicEvent
event EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
currEvents, b
newPop, Identifier
newId))
                      GenIO
gen
               else SimulationState b c -> IO (SimulationState b c)
forall (m :: * -> *) a. Monad m => a -> m a
return (SimulationState b c -> IO (SimulationState b c))
-> SimulationState b c -> IO (SimulationState b c)
forall a b. (a -> b) -> a -> b
$
                    (AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
forall b c.
(AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
SimulationState
                      ( AbsoluteTime
maxTime
                      , AbsoluteTime -> EpidemicEvent
StoppingTime AbsoluteTime
maxTime EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
currEvents
                      , b
currPop
                      , Identifier
currId)
           else SimulationState b c -> IO (SimulationState b c)
forall (m :: * -> *) a. Monad m => a -> m a
return (SimulationState b c -> IO (SimulationState b c))
-> SimulationState b c -> IO (SimulationState b c)
forall a b. (a -> b) -> a -> b
$
                (AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
forall b c.
(AbsoluteTime, [EpidemicEvent], b, Identifier)
-> SimulationState b c
SimulationState
                  ( AbsoluteTime
currTime
                  , AbsoluteTime -> EpidemicEvent
Extinction AbsoluteTime
currTime EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
currEvents
                  , b
currPop
                  , Identifier
currId)
     else SimulationState b c -> IO (SimulationState b c)
forall (m :: * -> *) a. Monad m => a -> m a
return (SimulationState b c -> IO (SimulationState b c))
-> (Maybe c -> SimulationState b c)
-> Maybe c
-> IO (SimulationState b c)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe c -> SimulationState b c
forall b c. Maybe c -> SimulationState b c
TerminatedSimulation (Maybe c -> IO (SimulationState b c))
-> Maybe c -> IO (SimulationState b c)
forall a b. (a -> b) -> a -> b
$ do TerminationHandler b -> Bool
_ [EpidemicEvent] -> c
termSummary <- Maybe (TerminationHandler b c)
maybeTermHandler
                                             c -> Maybe c
forall (m :: * -> *) a. Monad m => a -> m a
return (c -> Maybe c) -> c -> Maybe c
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent] -> c
termSummary [EpidemicEvent]
currEvents