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

module Epidemic.Types.Events
  ( EpidemicEvent(Infection, Removal, IndividualSample,
              PopulationSample, StoppingTime, Extinction)
  , popSampPeople
  , popSampSeq
  , popSampTime
  , indSampPerson
  , indSampSeq
  , indSampTime
  , EpidemicTree(Branch, Leaf, Shoot)
  , maybeEpidemicTree
  , isExtinctionOrStopping
  , isIndividualSample
  , derivedFrom
  ) where

import qualified Data.Aeson                as Json
import           Epidemic.Types.Population
import           Epidemic.Types.Time       (AbsoluteTime (..), TimeStamp (..))
import           GHC.Generics

-- | Events that can occur in an epidemic with their absolute time.
data EpidemicEvent
  = Infection AbsoluteTime Person Person -- ^ absolute time; infector; infectee
  | Removal AbsoluteTime Person
  | IndividualSample
      { EpidemicEvent -> AbsoluteTime
indSampTime   :: AbsoluteTime
      , EpidemicEvent -> Person
indSampPerson :: Person
      , EpidemicEvent -> Bool
indSampSeq    :: Bool
      }
  | PopulationSample
      { EpidemicEvent -> AbsoluteTime
popSampTime   :: AbsoluteTime
      , EpidemicEvent -> People
popSampPeople :: People
      , EpidemicEvent -> Bool
popSampSeq    :: Bool
      }
  | Extinction AbsoluteTime -- ^ epidemic went extinct
  | StoppingTime AbsoluteTime -- ^ the simulation reached the stopping time
  deriving (Int -> EpidemicEvent -> ShowS
[EpidemicEvent] -> ShowS
EpidemicEvent -> String
(Int -> EpidemicEvent -> ShowS)
-> (EpidemicEvent -> String)
-> ([EpidemicEvent] -> ShowS)
-> Show EpidemicEvent
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [EpidemicEvent] -> ShowS
$cshowList :: [EpidemicEvent] -> ShowS
show :: EpidemicEvent -> String
$cshow :: EpidemicEvent -> String
showsPrec :: Int -> EpidemicEvent -> ShowS
$cshowsPrec :: Int -> EpidemicEvent -> ShowS
Show, (forall x. EpidemicEvent -> Rep EpidemicEvent x)
-> (forall x. Rep EpidemicEvent x -> EpidemicEvent)
-> Generic EpidemicEvent
forall x. Rep EpidemicEvent x -> EpidemicEvent
forall x. EpidemicEvent -> Rep EpidemicEvent x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep EpidemicEvent x -> EpidemicEvent
$cfrom :: forall x. EpidemicEvent -> Rep EpidemicEvent x
Generic, EpidemicEvent -> EpidemicEvent -> Bool
(EpidemicEvent -> EpidemicEvent -> Bool)
-> (EpidemicEvent -> EpidemicEvent -> Bool) -> Eq EpidemicEvent
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: EpidemicEvent -> EpidemicEvent -> Bool
$c/= :: EpidemicEvent -> EpidemicEvent -> Bool
== :: EpidemicEvent -> EpidemicEvent -> Bool
$c== :: EpidemicEvent -> EpidemicEvent -> Bool
Eq)

instance Json.FromJSON EpidemicEvent

instance Json.ToJSON EpidemicEvent

instance TimeStamp EpidemicEvent where
  absTime :: EpidemicEvent -> AbsoluteTime
absTime EpidemicEvent
ee =
    case EpidemicEvent
ee of
      Infection AbsoluteTime
absT Person
_ Person
_    -> AbsoluteTime
absT
      Removal AbsoluteTime
absT Person
_        -> AbsoluteTime
absT
      IndividualSample {Bool
Person
AbsoluteTime
indSampSeq :: Bool
indSampPerson :: Person
indSampTime :: AbsoluteTime
indSampTime :: EpidemicEvent -> AbsoluteTime
indSampSeq :: EpidemicEvent -> Bool
indSampPerson :: EpidemicEvent -> Person
..} -> AbsoluteTime
indSampTime
      PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
popSampTime :: EpidemicEvent -> AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
..} -> AbsoluteTime
popSampTime
      StoppingTime AbsoluteTime
absT     -> AbsoluteTime
absT
      Extinction AbsoluteTime
absT       -> AbsoluteTime
absT

-- | Predicate for the event being an individual sample event.
isIndividualSample :: EpidemicEvent -> Bool
isIndividualSample :: EpidemicEvent -> Bool
isIndividualSample EpidemicEvent
ee =
  case EpidemicEvent
ee of
    IndividualSample {} -> Bool
True
    EpidemicEvent
_                   -> Bool
False

-- | Predicate for whether an @EpidemicEvent@ is one of the terminal events of
-- extinction or the stopping time having been reached.
isExtinctionOrStopping :: EpidemicEvent -> Bool
isExtinctionOrStopping :: EpidemicEvent -> Bool
isExtinctionOrStopping EpidemicEvent
e =
  case EpidemicEvent
e of
    Extinction {}   -> Bool
True
    StoppingTime {} -> Bool
True
    EpidemicEvent
_               -> Bool
False

-- | Epidemic Events are ordered based on which occurred first. Since
-- 'Extinction' and 'StoppingTime' events are there as placeholders they are
-- placed as the end of the order.
instance Ord EpidemicEvent where
  EpidemicEvent
e1 <= :: EpidemicEvent -> EpidemicEvent -> Bool
<= EpidemicEvent
e2 = EpidemicEvent -> AbsoluteTime
forall a. TimeStamp a => a -> AbsoluteTime
absTime EpidemicEvent
e1 AbsoluteTime -> AbsoluteTime -> Bool
forall a. Ord a => a -> a -> Bool
<= EpidemicEvent -> AbsoluteTime
forall a. TimeStamp a => a -> AbsoluteTime
absTime EpidemicEvent
e2

-- | The events that occurred as a result of the existance of the given person.
derivedFrom ::
     Person
  -> [EpidemicEvent] -- ^ ordered epidemic events
  -> [EpidemicEvent]
derivedFrom :: Person -> [EpidemicEvent] -> [EpidemicEvent]
derivedFrom Person
person = People -> [EpidemicEvent] -> [EpidemicEvent]
derivedFromPeople ([Person] -> People
asPeople [Person
person])

-- | The events that occurred as a result of the existance of a group of people
derivedFromPeople ::
     People
  -> [EpidemicEvent] -- ^ ordered epidemic events
  -> [EpidemicEvent]
derivedFromPeople :: People -> [EpidemicEvent] -> [EpidemicEvent]
derivedFromPeople People
_ [] = []
derivedFromPeople People
people (EpidemicEvent
e:[EpidemicEvent]
es) =
  case EpidemicEvent
e of
    Infection AbsoluteTime
_ Person
p1 Person
p2 ->
      if People -> Person -> Bool
includesPerson People
people Person
p1 Bool -> Bool -> Bool
|| People -> Person -> Bool
includesPerson People
people Person
p2
        then let people' :: People
people' = Person -> People -> People
addPerson Person
p2 (Person -> People -> People
addPerson Person
p1 People
people)
              in EpidemicEvent
e EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: People -> [EpidemicEvent] -> [EpidemicEvent]
derivedFromPeople People
people' [EpidemicEvent]
es
        else People -> [EpidemicEvent] -> [EpidemicEvent]
derivedFromPeople People
people [EpidemicEvent]
es
    Removal AbsoluteTime
_ Person
p ->
      let derivedEvents :: [EpidemicEvent]
derivedEvents = People -> [EpidemicEvent] -> [EpidemicEvent]
derivedFromPeople People
people [EpidemicEvent]
es
       in if People -> Person -> Bool
includesPerson People
people Person
p
            then EpidemicEvent
e EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
derivedEvents
            else [EpidemicEvent]
derivedEvents
    IndividualSample {Bool
Person
AbsoluteTime
indSampSeq :: Bool
indSampPerson :: Person
indSampTime :: AbsoluteTime
indSampTime :: EpidemicEvent -> AbsoluteTime
indSampSeq :: EpidemicEvent -> Bool
indSampPerson :: EpidemicEvent -> Person
..} ->
      let derivedEvents :: [EpidemicEvent]
derivedEvents = People -> [EpidemicEvent] -> [EpidemicEvent]
derivedFromPeople People
people [EpidemicEvent]
es
       in if People -> Person -> Bool
includesPerson People
people Person
indSampPerson
            then EpidemicEvent
e EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
derivedEvents
            else [EpidemicEvent]
derivedEvents
    PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
popSampTime :: EpidemicEvent -> AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
..} ->
      let derivedEvents :: [EpidemicEvent]
derivedEvents = People -> [EpidemicEvent] -> [EpidemicEvent]
derivedFromPeople People
people [EpidemicEvent]
es
       in if People -> People -> Bool
haveCommonPeople People
people People
popSampPeople
            then EpidemicEvent
e EpidemicEvent -> [EpidemicEvent] -> [EpidemicEvent]
forall a. a -> [a] -> [a]
: [EpidemicEvent]
derivedEvents
            else [EpidemicEvent]
derivedEvents
    Extinction {} -> People -> [EpidemicEvent] -> [EpidemicEvent]
derivedFromPeople People
people [EpidemicEvent]
es
    StoppingTime {} -> People -> [EpidemicEvent] -> [EpidemicEvent]
derivedFromPeople People
people [EpidemicEvent]
es

-- | The whole transmission tree including the unobserved leaves. Lineages that
-- are still extant are modelled as /shoots/ and contain a 'Person' as their
-- data rather than an event.
data EpidemicTree
  = Branch EpidemicEvent EpidemicTree EpidemicTree
  | Leaf EpidemicEvent
  | Shoot Person
  deriving (Int -> EpidemicTree -> ShowS
[EpidemicTree] -> ShowS
EpidemicTree -> String
(Int -> EpidemicTree -> ShowS)
-> (EpidemicTree -> String)
-> ([EpidemicTree] -> ShowS)
-> Show EpidemicTree
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [EpidemicTree] -> ShowS
$cshowList :: [EpidemicTree] -> ShowS
show :: EpidemicTree -> String
$cshow :: EpidemicTree -> String
showsPrec :: Int -> EpidemicTree -> ShowS
$cshowsPrec :: Int -> EpidemicTree -> ShowS
Show, EpidemicTree -> EpidemicTree -> Bool
(EpidemicTree -> EpidemicTree -> Bool)
-> (EpidemicTree -> EpidemicTree -> Bool) -> Eq EpidemicTree
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: EpidemicTree -> EpidemicTree -> Bool
$c/= :: EpidemicTree -> EpidemicTree -> Bool
== :: EpidemicTree -> EpidemicTree -> Bool
$c== :: EpidemicTree -> EpidemicTree -> Bool
Eq)

-- | If possible return an 'EpidemicTree' describing the /sorted/ list of
-- 'EpidemicEvent'.
maybeEpidemicTree ::
     [EpidemicEvent] -- ^ ordered epidemic events
  -> Either String EpidemicTree
maybeEpidemicTree :: [EpidemicEvent] -> Either String EpidemicTree
maybeEpidemicTree [] =
  String -> Either String EpidemicTree
forall a b. a -> Either a b
Left String
"There are no EpidemicEvent values to construct a tree with."
maybeEpidemicTree [EpidemicEvent
e] =
  case EpidemicEvent
e of
    Infection AbsoluteTime
_ Person
p1 Person
p2 -> EpidemicTree -> Either String EpidemicTree
forall a b. b -> Either a b
Right (EpidemicEvent -> EpidemicTree -> EpidemicTree -> EpidemicTree
Branch EpidemicEvent
e (Person -> EpidemicTree
Shoot Person
p1) (Person -> EpidemicTree
Shoot Person
p2))
    Removal {} -> EpidemicTree -> Either String EpidemicTree
forall a b. b -> Either a b
Right (EpidemicEvent -> EpidemicTree
Leaf EpidemicEvent
e)
    IndividualSample {} -> EpidemicTree -> Either String EpidemicTree
forall a b. b -> Either a b
Right (EpidemicEvent -> EpidemicTree
Leaf EpidemicEvent
e)
    PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
popSampTime :: EpidemicEvent -> AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
..} ->
      if People -> Bool
nullPeople People
popSampPeople
        then String -> Either String EpidemicTree
forall a b. a -> Either a b
Left String
"The last event is a PopulationSample with no people sampled"
        else EpidemicTree -> Either String EpidemicTree
forall a b. b -> Either a b
Right (EpidemicEvent -> EpidemicTree
Leaf EpidemicEvent
e)
    Extinction {} ->
      String -> Either String EpidemicTree
forall a b. a -> Either a b
Left String
"Extinction event encountered. It should have been removed"
    StoppingTime {} ->
      String -> Either String EpidemicTree
forall a b. a -> Either a b
Left String
"Stopping time encountered. It should have been removed"
maybeEpidemicTree (EpidemicEvent
e:[EpidemicEvent]
es) =
  case EpidemicEvent
e of
    Infection AbsoluteTime
_ Person
p1 Person
p2 ->
      let infectorEvents :: [EpidemicEvent]
infectorEvents = Person -> [EpidemicEvent] -> [EpidemicEvent]
derivedFrom Person
p1 [EpidemicEvent]
es
          infecteeEvents :: [EpidemicEvent]
infecteeEvents = Person -> [EpidemicEvent] -> [EpidemicEvent]
derivedFrom Person
p2 [EpidemicEvent]
es
       in do EpidemicTree
leftTree <-
               if [EpidemicEvent] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [EpidemicEvent]
infectorEvents
                 then EpidemicTree -> Either String EpidemicTree
forall a b. b -> Either a b
Right (Person -> EpidemicTree
Shoot Person
p1)
                 else [EpidemicEvent] -> Either String EpidemicTree
maybeEpidemicTree [EpidemicEvent]
infectorEvents
             EpidemicTree
rightTree <-
               if [EpidemicEvent] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [EpidemicEvent]
infecteeEvents
                 then EpidemicTree -> Either String EpidemicTree
forall a b. b -> Either a b
Right (Person -> EpidemicTree
Shoot Person
p2)
                 else [EpidemicEvent] -> Either String EpidemicTree
maybeEpidemicTree [EpidemicEvent]
infecteeEvents
             EpidemicTree -> Either String EpidemicTree
forall (m :: * -> *) a. Monad m => a -> m a
return (EpidemicTree -> Either String EpidemicTree)
-> EpidemicTree -> Either String EpidemicTree
forall a b. (a -> b) -> a -> b
$ EpidemicEvent -> EpidemicTree -> EpidemicTree -> EpidemicTree
Branch EpidemicEvent
e EpidemicTree
leftTree EpidemicTree
rightTree
    Removal {} -> EpidemicTree -> Either String EpidemicTree
forall a b. b -> Either a b
Right (EpidemicEvent -> EpidemicTree
Leaf EpidemicEvent
e)
    IndividualSample {} -> EpidemicTree -> Either String EpidemicTree
forall a b. b -> Either a b
Right (EpidemicEvent -> EpidemicTree
Leaf EpidemicEvent
e)
    PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
popSampTime :: EpidemicEvent -> AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
..} ->
      if People -> Bool
nullPeople People
popSampPeople
        then [EpidemicEvent] -> Either String EpidemicTree
maybeEpidemicTree [EpidemicEvent]
es
        else EpidemicTree -> Either String EpidemicTree
forall a b. b -> Either a b
Right (EpidemicEvent -> EpidemicTree
Leaf EpidemicEvent
e)
    Extinction {} ->
      String -> Either String EpidemicTree
forall a b. a -> Either a b
Left String
"Extinction event encountered. It should have been removed"
    StoppingTime {} ->
      String -> Either String EpidemicTree
forall a b. a -> Either a b
Left String
"Stopping time encountered. It should have been removed"