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

module Epidemic where

import Control.Monad
import qualified Data.ByteString as B
import Data.ByteString.Internal (c2w)
import Data.List (nub)
import Data.Maybe (fromJust, isJust, isNothing)
import qualified Data.Vector as V
import Data.Word
import Epidemic.Types.Events
import Epidemic.Types.Parameter
import Epidemic.Types.Population
import Epidemic.Types.Simulation
  ( SimulationConfiguration(..)
  , SimulationRandEvent(..)
  , SimulationState(..)
  )
import Epidemic.Types.Time (AbsoluteTime(..), Timed(..), diracDeltaValue, nextTime)
import GHC.Generics (Generic)
import System.Random.MWC

-- | The number of people added or removed in an event.
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
    EpidemicEvent
StoppingTime -> 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.
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
    EpidemicEvent
Extinction -> []
    EpidemicEvent
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 and return a @SimulationState@ which holds the history
-- of the simulation.
allEvents ::
     (ModelParameters a b, Population b)
  => SimulationRandEvent a b
  -> a
  -> AbsoluteTime
  -> Maybe (b -> Bool) -- ^ predicate for a valid population
  -> SimulationState b
  -> GenIO
  -> IO (SimulationState b)
allEvents :: SimulationRandEvent a b
-> a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEvents SimulationRandEvent a b
_ a
_ AbsoluteTime
_ Maybe (b -> Bool)
_ SimulationState b
TerminatedSimulation GenIO
_ = SimulationState b -> IO (SimulationState b)
forall (m :: * -> *) a. Monad m => a -> m a
return SimulationState b
forall b. SimulationState b
TerminatedSimulation
allEvents simRandEvent :: SimulationRandEvent a b
simRandEvent@(SimulationRandEvent a
-> AbsoluteTime
-> b
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, b, Identifier)
randEvent) a
modelParams AbsoluteTime
maxTime Maybe (b -> Bool)
maybePopPredicate (SimulationState (AbsoluteTime
currTime, [EpidemicEvent]
currEvents, b
currPop, Identifier
currId)) GenIO
gen =
  if Maybe (b -> Bool) -> Bool
forall a. Maybe a -> Bool
isNothing Maybe (b -> Bool)
maybePopPredicate Bool -> Bool -> Bool
||
     (Maybe (b -> Bool) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (b -> Bool)
maybePopPredicate Bool -> Bool -> Bool
&& Maybe (b -> Bool) -> b -> Bool
forall a. HasCallStack => Maybe a -> a
fromJust Maybe (b -> Bool)
maybePopPredicate 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 (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
forall a b.
(ModelParameters a b, Population b) =>
SimulationRandEvent a b
-> a
-> AbsoluteTime
-> Maybe (b -> Bool)
-> SimulationState b
-> GenIO
-> IO (SimulationState b)
allEvents
                      SimulationRandEvent a b
simRandEvent
                      a
modelParams
                      AbsoluteTime
maxTime
                      Maybe (b -> Bool)
maybePopPredicate
                      ((AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
forall b.
(AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
SimulationState
                         (AbsoluteTime
newTime, EpidemicEvent
event EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
currEvents, b
newPop, Identifier
newId))
                      GenIO
gen
               else SimulationState b -> IO (SimulationState b)
forall (m :: * -> *) a. Monad m => a -> m a
return (SimulationState b -> IO (SimulationState b))
-> SimulationState b -> IO (SimulationState b)
forall a b. (a -> b) -> a -> b
$
                    (AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
forall b.
(AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
SimulationState
                      (AbsoluteTime
maxTime, EpidemicEvent
StoppingTime EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
currEvents, b
currPop, Identifier
currId)
           else SimulationState b -> IO (SimulationState b)
forall (m :: * -> *) a. Monad m => a -> m a
return (SimulationState b -> IO (SimulationState b))
-> SimulationState b -> IO (SimulationState b)
forall a b. (a -> b) -> a -> b
$
                (AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
forall b.
(AbsoluteTime, [EpidemicEvent], b, Identifier) -> SimulationState b
SimulationState
                  (AbsoluteTime
currTime, EpidemicEvent
Extinction EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
currEvents, b
currPop, Identifier
currId)
    else SimulationState b -> IO (SimulationState b)
forall (m :: * -> *) a. Monad m => a -> m a
return SimulationState b
forall b. SimulationState b
TerminatedSimulation