{-# LANGUAGE MultiParamTypeClasses #-}

module Epidemic.Model.InhomogeneousBDS
  ( configuration
  , randomEvent
  , inhomBDSRates
  , InhomBDSRates(..)
  , InhomBDSPop(..)
  ) where

import Epidemic.Types.Time
  ( AbsoluteTime(..)
  , Timed(..)
  , TimeDelta(..)
  , allTimes
  , asTimed
  , cadlagValue
  )
import Data.Maybe (fromJust)
import qualified Data.Vector as V
import Epidemic.Types.Events
  ( EpidemicEvent(..)
  )
import Epidemic.Types.Parameter
import Epidemic.Types.Population
import Epidemic.Types.Simulation
  ( SimulationConfiguration(..)
  , SimulationRandEvent(..), TerminationHandler(..)
  )
import Epidemic.Utility
import System.Random.MWC
import System.Random.MWC.Distributions (categorical)

data InhomBDSRates =
  InhomBDSRates (Timed Rate) Rate Rate

data InhomBDSPop =
  InhomBDSPop People
  deriving (Int -> InhomBDSPop -> ShowS
[InhomBDSPop] -> ShowS
InhomBDSPop -> String
(Int -> InhomBDSPop -> ShowS)
-> (InhomBDSPop -> String)
-> ([InhomBDSPop] -> ShowS)
-> Show InhomBDSPop
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [InhomBDSPop] -> ShowS
$cshowList :: [InhomBDSPop] -> ShowS
show :: InhomBDSPop -> String
$cshow :: InhomBDSPop -> String
showsPrec :: Int -> InhomBDSPop -> ShowS
$cshowsPrec :: Int -> InhomBDSPop -> ShowS
Show)

instance ModelParameters InhomBDSRates InhomBDSPop where
  rNaught :: InhomBDSPop -> InhomBDSRates -> AbsoluteTime -> Maybe Double
rNaught InhomBDSPop
_ (InhomBDSRates Timed Double
timedBirthRate Double
deathRate Double
sampleRate) AbsoluteTime
time =
    let birthRate :: Maybe Double
birthRate = Timed Double -> AbsoluteTime -> Maybe Double
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Double
timedBirthRate AbsoluteTime
time
     in (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
deathRate Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
sampleRate)) (Double -> Double) -> Maybe Double -> Maybe Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Double
birthRate
  eventRate :: InhomBDSPop -> InhomBDSRates -> AbsoluteTime -> Maybe Double
eventRate InhomBDSPop
_ (InhomBDSRates Timed Double
timedBirthRate Double
deathRate Double
sampleRate) AbsoluteTime
time =
    let birthRate :: Maybe Double
birthRate = Timed Double -> AbsoluteTime -> Maybe Double
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Double
timedBirthRate AbsoluteTime
time
     in (Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
deathRate Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
sampleRate)) (Double -> Double) -> Maybe Double -> Maybe Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Double
birthRate
  birthProb :: InhomBDSPop -> InhomBDSRates -> AbsoluteTime -> Maybe Double
birthProb InhomBDSPop
_ (InhomBDSRates Timed Double
timedBirthRate Double
deathRate Double
sampleRate) AbsoluteTime
time =
    (\Double
br -> Double
br Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
br Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
deathRate Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
sampleRate)) (Double -> Double) -> Maybe Double -> Maybe Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
    Timed Double -> AbsoluteTime -> Maybe Double
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Double
timedBirthRate AbsoluteTime
time
  eventWeights :: InhomBDSPop
-> InhomBDSRates -> AbsoluteTime -> Maybe (Vector Double)
eventWeights InhomBDSPop
_ (InhomBDSRates Timed Double
timedBirthRate Double
deathRate Double
sampleRate) AbsoluteTime
time =
    Vector Double -> Maybe (Vector Double)
forall a. a -> Maybe a
Just (Vector Double -> Maybe (Vector Double))
-> Vector Double -> Maybe (Vector Double)
forall a b. (a -> b) -> a -> b
$ [Double] -> Vector Double
forall a. [a] -> Vector a
V.fromList [Maybe Double -> Double
forall a. HasCallStack => Maybe a -> a
fromJust (Timed Double -> AbsoluteTime -> Maybe Double
forall a. Timed a -> AbsoluteTime -> Maybe a
cadlagValue Timed Double
timedBirthRate AbsoluteTime
time), Double
deathRate, Double
sampleRate]

instance Population InhomBDSPop where
  susceptiblePeople :: InhomBDSPop -> Maybe People
susceptiblePeople InhomBDSPop
_ = Maybe People
forall a. Maybe a
Nothing
  infectiousPeople :: InhomBDSPop -> Maybe People
infectiousPeople (InhomBDSPop People
people) = People -> Maybe People
forall a. a -> Maybe a
Just People
people
  removedPeople :: InhomBDSPop -> Maybe People
removedPeople InhomBDSPop
_ = Maybe People
forall a. Maybe a
Nothing
  isInfected :: InhomBDSPop -> Bool
isInfected (InhomBDSPop People
people) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ People -> Bool
nullPeople People
people

-- | Return a BDS-process parameters object
--
-- Note that this requires that the rates are all positive, if they are not it
-- will return @Nothing@.
inhomBDSRates ::
     Timed Rate -- ^ birth rate
  -> Rate -- ^ death rate
  -> Rate -- ^ sample rate
  -> Maybe InhomBDSRates
inhomBDSRates :: Timed Double -> Double -> Double -> Maybe InhomBDSRates
inhomBDSRates timedBirthRate :: Timed Double
timedBirthRate@(Timed [(AbsoluteTime, Double)]
tBrPairs) Double
deathRate Double
sampleRate
  | ((AbsoluteTime, Double) -> Bool)
-> [(AbsoluteTime, Double)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(AbsoluteTime, Double)
x -> Double
0 Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< (AbsoluteTime, Double) -> Double
forall a b. (a, b) -> b
snd (AbsoluteTime, Double)
x) [(AbsoluteTime, Double)]
tBrPairs Bool -> Bool -> Bool
&& Double
deathRate Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 Bool -> Bool -> Bool
&& Double
sampleRate Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 =
    InhomBDSRates -> Maybe InhomBDSRates
forall a. a -> Maybe a
Just (InhomBDSRates -> Maybe InhomBDSRates)
-> InhomBDSRates -> Maybe InhomBDSRates
forall a b. (a -> b) -> a -> b
$ Timed Double -> Double -> Double -> InhomBDSRates
InhomBDSRates Timed Double
timedBirthRate Double
deathRate Double
sampleRate
  | Bool
otherwise = Maybe InhomBDSRates
forall a. Maybe a
Nothing

-- | Configuration of a inhomogeneous birth-death-sampling simulation.
--
-- Note that this requires that the timed rates are all positive, if they are
-- not it will return @Nothing@ which can lead to cryptic bugs.
configuration ::
     TimeDelta -- ^ Duration of the simulation after starting at time 0.
  -> Bool -- ^ condition upon at least two sequenced samples.
  -> Maybe (InhomBDSPop -> Bool, [EpidemicEvent] -> s) -- ^ values for termination handling.
  -> ([(AbsoluteTime, Rate)], Rate, Rate) -- ^ Birth, Death and Sampling rates
  -> Maybe (SimulationConfiguration InhomBDSRates InhomBDSPop s)
configuration :: TimeDelta
-> Bool
-> Maybe (InhomBDSPop -> Bool, [EpidemicEvent] -> s)
-> ([(AbsoluteTime, Double)], Double, Double)
-> Maybe (SimulationConfiguration InhomBDSRates InhomBDSPop s)
configuration TimeDelta
maxTime Bool
atLeastCherry Maybe (InhomBDSPop -> Bool, [EpidemicEvent] -> s)
maybeTHFuncs ([(AbsoluteTime, Double)]
tBrPairs, Double
deathRate, Double
sampleRate) =
  let (Person
seedPerson, Identifier
newId) = Identifier -> (Person, Identifier)
newPerson Identifier
initialIdentifier
      bdsPop :: InhomBDSPop
bdsPop = People -> InhomBDSPop
InhomBDSPop (Vector Person -> People
People (Vector Person -> People) -> Vector Person -> People
forall a b. (a -> b) -> a -> b
$ Person -> Vector Person
forall a. a -> Vector a
V.singleton Person
seedPerson)
      termHandler :: Maybe (TerminationHandler InhomBDSPop s)
termHandler = do (InhomBDSPop -> Bool
f1, [EpidemicEvent] -> s
f2) <- Maybe (InhomBDSPop -> Bool, [EpidemicEvent] -> s)
maybeTHFuncs
                       TerminationHandler InhomBDSPop s
-> Maybe (TerminationHandler InhomBDSPop s)
forall (m :: * -> *) a. Monad m => a -> m a
return (TerminationHandler InhomBDSPop s
 -> Maybe (TerminationHandler InhomBDSPop s))
-> TerminationHandler InhomBDSPop s
-> Maybe (TerminationHandler InhomBDSPop s)
forall a b. (a -> b) -> a -> b
$ (InhomBDSPop -> Bool)
-> ([EpidemicEvent] -> s) -> TerminationHandler InhomBDSPop s
forall b c.
Population b =>
(b -> Bool) -> ([EpidemicEvent] -> c) -> TerminationHandler b c
TerminationHandler InhomBDSPop -> Bool
f1 [EpidemicEvent] -> s
f2
   in do Timed Double
timedBirthRate <- [(AbsoluteTime, Double)] -> Maybe (Timed Double)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed [(AbsoluteTime, Double)]
tBrPairs
         InhomBDSRates
maybeIBDSRates <- Timed Double -> Double -> Double -> Maybe InhomBDSRates
inhomBDSRates Timed Double
timedBirthRate Double
deathRate Double
sampleRate
         if TimeDelta
maxTime TimeDelta -> TimeDelta -> Bool
forall a. Ord a => a -> a -> Bool
> Double -> TimeDelta
TimeDelta Double
0
           then SimulationConfiguration InhomBDSRates InhomBDSPop s
-> Maybe (SimulationConfiguration InhomBDSRates InhomBDSPop s)
forall a. a -> Maybe a
Just
                  (InhomBDSRates
-> InhomBDSPop
-> Identifier
-> AbsoluteTime
-> TimeDelta
-> Maybe (TerminationHandler InhomBDSPop s)
-> Bool
-> SimulationConfiguration InhomBDSRates InhomBDSPop s
forall r p s.
r
-> p
-> Identifier
-> AbsoluteTime
-> TimeDelta
-> Maybe (TerminationHandler p s)
-> Bool
-> SimulationConfiguration r p s
SimulationConfiguration
                     InhomBDSRates
maybeIBDSRates
                     InhomBDSPop
bdsPop
                     Identifier
newId
                     (Double -> AbsoluteTime
AbsoluteTime Double
0)
                     TimeDelta
maxTime
                     Maybe (TerminationHandler InhomBDSPop s)
termHandler
                     Bool
atLeastCherry)
           else Maybe (SimulationConfiguration InhomBDSRates InhomBDSPop s)
forall a. Maybe a
Nothing

randomEvent :: SimulationRandEvent InhomBDSRates InhomBDSPop
randomEvent :: SimulationRandEvent InhomBDSRates InhomBDSPop
randomEvent = (InhomBDSRates
 -> AbsoluteTime
 -> InhomBDSPop
 -> Identifier
 -> GenIO
 -> IO (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier))
-> SimulationRandEvent InhomBDSRates InhomBDSPop
forall a b.
(ModelParameters a b, Population b) =>
(a
 -> AbsoluteTime
 -> b
 -> Identifier
 -> GenIO
 -> IO (AbsoluteTime, EpidemicEvent, b, Identifier))
-> SimulationRandEvent a b
SimulationRandEvent InhomBDSRates
-> AbsoluteTime
-> InhomBDSPop
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier)
randomEvent'

-- | A random event and the state afterwards
randomEvent' ::
     InhomBDSRates -- ^ model parameters
  -> AbsoluteTime -- ^ the current time
  -> InhomBDSPop -- ^ the population
  -> Identifier -- ^ current identifier
  -> GenIO -- ^ PRNG
  -> IO (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier)
randomEvent' :: InhomBDSRates
-> AbsoluteTime
-> InhomBDSPop
-> Identifier
-> GenIO
-> IO (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier)
randomEvent' inhomRates :: InhomBDSRates
inhomRates@(InhomBDSRates Timed Double
brts Double
_ Double
_) AbsoluteTime
currTime pop :: InhomBDSPop
pop@(InhomBDSPop People
people) Identifier
currId GenIO
gen =
  let popSize :: Double
popSize = 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
      --weightVecFunc :: AbsoluteTime -> Maybe (Vector Double)
      weightVecFunc :: AbsoluteTime -> Maybe (Vector Double)
weightVecFunc = InhomBDSPop
-> InhomBDSRates -> AbsoluteTime -> Maybe (Vector Double)
forall a p.
ModelParameters a p =>
p -> a -> AbsoluteTime -> Maybe (Vector Double)
eventWeights InhomBDSPop
pop InhomBDSRates
inhomRates
      -- we need a new step function to account for the population size.
      (Just Timed Double
stepFunction) =
        [(AbsoluteTime, Double)] -> Maybe (Timed Double)
forall a. Num a => [(AbsoluteTime, a)] -> Maybe (Timed a)
asTimed
          [ (AbsoluteTime
t, Double
popSize Double -> Double -> Double
forall a. Num a => a -> a -> a
* Maybe Double -> Double
forall a. HasCallStack => Maybe a -> a
fromJust (InhomBDSPop -> InhomBDSRates -> AbsoluteTime -> Maybe Double
forall a p.
ModelParameters a p =>
p -> a -> AbsoluteTime -> Maybe Double
eventRate InhomBDSPop
pop InhomBDSRates
inhomRates AbsoluteTime
t))
          | AbsoluteTime
t <- Timed Double -> [AbsoluteTime]
forall a. Timed a -> [AbsoluteTime]
allTimes Timed Double
brts
          ]
   in do (Just AbsoluteTime
newEventTime) <- Timed Double -> AbsoluteTime -> GenIO -> IO (Maybe AbsoluteTime)
forall (m :: * -> *).
PrimMonad m =>
Timed Double
-> AbsoluteTime -> Gen (PrimState m) -> m (Maybe AbsoluteTime)
inhomExponential Timed Double
stepFunction AbsoluteTime
currTime GenIO
gen
         Int
eventIx <- Vector Double -> Gen RealWorld -> IO Int
forall g (m :: * -> *) (v :: * -> *).
(StatefulGen g m, Vector v Double) =>
v Double -> g -> m Int
categorical (Maybe (Vector Double) -> Vector Double
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Vector Double) -> Vector Double)
-> Maybe (Vector Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ AbsoluteTime -> Maybe (Vector Double)
weightVecFunc AbsoluteTime
newEventTime) Gen RealWorld
GenIO
gen
         (Person
selectedPerson, People
unselectedPeople) <- People -> GenIO -> IO (Person, People)
randomPerson People
people GenIO
gen
         (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier)
-> IO (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier)
forall (m :: * -> *) a. Monad m => a -> m a
return ((AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier)
 -> IO (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier))
-> (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier)
-> IO (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier)
forall a b. (a -> b) -> a -> b
$
           case Int
eventIx of
             Int
0 ->
               ( AbsoluteTime
newEventTime
               , AbsoluteTime -> Person -> Person -> EpidemicEvent
Infection AbsoluteTime
newEventTime Person
selectedPerson Person
birthedPerson
               , People -> InhomBDSPop
InhomBDSPop (Person -> People -> People
addPerson Person
birthedPerson People
people)
               , Identifier
newId)
               where (Person
birthedPerson, Identifier
newId) = Identifier -> (Person, Identifier)
newPerson Identifier
currId
             Int
1 ->
               ( AbsoluteTime
newEventTime
               , AbsoluteTime -> Person -> EpidemicEvent
Removal AbsoluteTime
newEventTime Person
selectedPerson
               , People -> InhomBDSPop
InhomBDSPop People
unselectedPeople
               , Identifier
currId)
             Int
2 ->
               ( AbsoluteTime
newEventTime
               , AbsoluteTime -> Person -> Bool -> EpidemicEvent
IndividualSample AbsoluteTime
newEventTime Person
selectedPerson Bool
True
               , People -> InhomBDSPop
InhomBDSPop People
unselectedPeople
               , Identifier
currId)
             Int
_ -> String -> (AbsoluteTime, EpidemicEvent, InhomBDSPop, Identifier)
forall a. HasCallStack => String -> a
error String
"no birth-death-sampling event selected."