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

module Epidemic.Types.Observations
  ( Observation(..)
  , ReconstructedTree(..)
  , maybeReconstructedTree
  , PointProcessEvents(..)
  , pointProcessEvents
  , reconstructedTreeEvents
  , observedEvents
  ) where

import Control.Monad (liftM)
import qualified Data.Aeson as Json
import qualified Data.ByteString.Builder as BBuilder
import qualified Data.List as List
import qualified Data.Vector as V
import Epidemic.Types.Events
  ( EpidemicEvent(..)
  , EpidemicTree(..)
  , maybeEpidemicTree
  )
import Epidemic.Types.Time (TimeDelta(..), timeDelta)
import Epidemic.Types.Population (People(..), personByteString)
import GHC.Generics

-- | A wrapper for an 'EpidemicEvent' to indicate that this is an even that was
-- observed rather than just an event of the epidemic process.
newtype Observation =
  Observation EpidemicEvent
  deriving (Int -> Observation -> ShowS
[Observation] -> ShowS
Observation -> String
(Int -> Observation -> ShowS)
-> (Observation -> String)
-> ([Observation] -> ShowS)
-> Show Observation
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Observation] -> ShowS
$cshowList :: [Observation] -> ShowS
show :: Observation -> String
$cshow :: Observation -> String
showsPrec :: Int -> Observation -> ShowS
$cshowsPrec :: Int -> Observation -> ShowS
Show, Eq Observation
Eq Observation
-> (Observation -> Observation -> Ordering)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool)
-> (Observation -> Observation -> Observation)
-> (Observation -> Observation -> Observation)
-> Ord Observation
Observation -> Observation -> Bool
Observation -> Observation -> Ordering
Observation -> Observation -> Observation
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Observation -> Observation -> Observation
$cmin :: Observation -> Observation -> Observation
max :: Observation -> Observation -> Observation
$cmax :: Observation -> Observation -> Observation
>= :: Observation -> Observation -> Bool
$c>= :: Observation -> Observation -> Bool
> :: Observation -> Observation -> Bool
$c> :: Observation -> Observation -> Bool
<= :: Observation -> Observation -> Bool
$c<= :: Observation -> Observation -> Bool
< :: Observation -> Observation -> Bool
$c< :: Observation -> Observation -> Bool
compare :: Observation -> Observation -> Ordering
$ccompare :: Observation -> Observation -> Ordering
$cp1Ord :: Eq Observation
Ord, Observation -> Observation -> Bool
(Observation -> Observation -> Bool)
-> (Observation -> Observation -> Bool) -> Eq Observation
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Observation -> Observation -> Bool
$c/= :: Observation -> Observation -> Bool
== :: Observation -> Observation -> Bool
$c== :: Observation -> Observation -> Bool
Eq, (forall x. Observation -> Rep Observation x)
-> (forall x. Rep Observation x -> Observation)
-> Generic Observation
forall x. Rep Observation x -> Observation
forall x. Observation -> Rep Observation x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Observation x -> Observation
$cfrom :: forall x. Observation -> Rep Observation x
Generic)

instance Json.FromJSON Observation

instance Json.ToJSON Observation

-- | A representation of the events that can be observed in an epidemic but
-- which are not included in the reconstructed tree, ie the unsequenced
-- observations.
newtype PointProcessEvents =
  PointProcessEvents [Observation]

-- | Extract the events from an epidemic tree which are observed but not part of
-- the reconstructed tree, ie the ones that are not sequenced.
pointProcessEvents :: EpidemicTree -> PointProcessEvents
pointProcessEvents :: EpidemicTree -> PointProcessEvents
pointProcessEvents Shoot {} = [Observation] -> PointProcessEvents
PointProcessEvents []
pointProcessEvents (Leaf 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
..} -> [Observation] -> PointProcessEvents
PointProcessEvents ([Observation] -> PointProcessEvents)
-> [Observation] -> PointProcessEvents
forall a b. (a -> b) -> a -> b
$ if Bool -> Bool
not Bool
indSampSeq then [EpidemicEvent -> Observation
Observation EpidemicEvent
e] else []
    PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
..} -> [Observation] -> PointProcessEvents
PointProcessEvents ([Observation] -> PointProcessEvents)
-> [Observation] -> PointProcessEvents
forall a b. (a -> b) -> a -> b
$ if Bool -> Bool
not Bool
popSampSeq then [EpidemicEvent -> Observation
Observation EpidemicEvent
e] else []
    EpidemicEvent
_ -> [Observation] -> PointProcessEvents
PointProcessEvents []
pointProcessEvents (Branch EpidemicEvent
_ EpidemicTree
lt EpidemicTree
rt) =
  let (PointProcessEvents [Observation]
lEs) = EpidemicTree -> PointProcessEvents
pointProcessEvents EpidemicTree
lt
      (PointProcessEvents [Observation]
rEs) = EpidemicTree -> PointProcessEvents
pointProcessEvents EpidemicTree
rt
      allEs :: [Observation]
allEs = [Observation] -> [Observation]
forall a. Ord a => [a] -> [a]
List.sort ([Observation] -> [Observation]) -> [Observation] -> [Observation]
forall a b. (a -> b) -> a -> b
$ [Observation]
lEs [Observation] -> [Observation] -> [Observation]
forall a. [a] -> [a] -> [a]
++ [Observation]
rEs
   in [Observation] -> PointProcessEvents
PointProcessEvents [Observation]
allEs

-- | A representation of the reconstructed tree, ie the tree where the leaves
-- correspond to sequenced observations.
data ReconstructedTree
  = RBranch Observation ReconstructedTree ReconstructedTree
  | RLeaf Observation
  deriving (Int -> ReconstructedTree -> ShowS
[ReconstructedTree] -> ShowS
ReconstructedTree -> String
(Int -> ReconstructedTree -> ShowS)
-> (ReconstructedTree -> String)
-> ([ReconstructedTree] -> ShowS)
-> Show ReconstructedTree
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ReconstructedTree] -> ShowS
$cshowList :: [ReconstructedTree] -> ShowS
show :: ReconstructedTree -> String
$cshow :: ReconstructedTree -> String
showsPrec :: Int -> ReconstructedTree -> ShowS
$cshowsPrec :: Int -> ReconstructedTree -> ShowS
Show, ReconstructedTree -> ReconstructedTree -> Bool
(ReconstructedTree -> ReconstructedTree -> Bool)
-> (ReconstructedTree -> ReconstructedTree -> Bool)
-> Eq ReconstructedTree
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ReconstructedTree -> ReconstructedTree -> Bool
$c/= :: ReconstructedTree -> ReconstructedTree -> Bool
== :: ReconstructedTree -> ReconstructedTree -> Bool
$c== :: ReconstructedTree -> ReconstructedTree -> Bool
Eq)

-- | The reconstructed phylogeny obtained by pruning an 'EpidemicTree' which
-- contains represents the transmission tree of the epidemic. In the case where
-- there are no sequenced samples in the epidemic then there is no tree to
-- reconstruct which is why this function is in the either monad.
maybeReconstructedTree :: EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree :: EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree Shoot {} = String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"EpidemicTree is only a Shoot"
maybeReconstructedTree (Leaf EpidemicEvent
e) =
  case EpidemicEvent
e of
    IndividualSample {Bool
Person
AbsoluteTime
indSampSeq :: Bool
indSampPerson :: Person
indSampTime :: AbsoluteTime
indSampSeq :: EpidemicEvent -> Bool
indSampPerson :: EpidemicEvent -> Person
indSampTime :: EpidemicEvent -> AbsoluteTime
..} -> if Bool
indSampSeq
                             then ReconstructedTree -> Either String ReconstructedTree
forall a b. b -> Either a b
Right (ReconstructedTree -> Either String ReconstructedTree)
-> ReconstructedTree -> Either String ReconstructedTree
forall a b. (a -> b) -> a -> b
$ Observation -> ReconstructedTree
RLeaf (EpidemicEvent -> Observation
Observation EpidemicEvent
e)
                             else String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"Leaf with non-sequenced event individual sample"
    PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
..} -> if Bool
popSampSeq
                             then ReconstructedTree -> Either String ReconstructedTree
forall a b. b -> Either a b
Right (ReconstructedTree -> Either String ReconstructedTree)
-> ReconstructedTree -> Either String ReconstructedTree
forall a b. (a -> b) -> a -> b
$ Observation -> ReconstructedTree
RLeaf (EpidemicEvent -> Observation
Observation EpidemicEvent
e)
                             else String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"Leaf with non-sequenced event population sample"
    EpidemicEvent
_ -> String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"Bad leaf in the EpidemicTree"
maybeReconstructedTree (Branch e :: EpidemicEvent
e@Infection {} EpidemicTree
lt EpidemicTree
rt)
  | EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
lt Bool -> Bool -> Bool
&& EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
rt = do
    ReconstructedTree
rlt <- EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree EpidemicTree
lt
    ReconstructedTree
rrt <- EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree EpidemicTree
rt
    ReconstructedTree -> Either String ReconstructedTree
forall a b. b -> Either a b
Right (ReconstructedTree -> Either String ReconstructedTree)
-> ReconstructedTree -> Either String ReconstructedTree
forall a b. (a -> b) -> a -> b
$ Observation
-> ReconstructedTree -> ReconstructedTree -> ReconstructedTree
RBranch (EpidemicEvent -> Observation
Observation EpidemicEvent
e) ReconstructedTree
rlt ReconstructedTree
rrt
  | EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
lt = EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree EpidemicTree
lt
  | EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
rt = EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree EpidemicTree
rt
  | Bool
otherwise = String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"Neither subtree has a sequenced leaf"
maybeReconstructedTree Branch {} = String -> Either String ReconstructedTree
forall a b. a -> Either a b
Left String
"EpidemicTree is a bad branch"

-- | Predicate for whether an 'EpidemicTree' has any leaf which corresponds to a
-- sequenced observation and hence should be included in a @ReconstructedTree@.
hasSequencedLeaf :: EpidemicTree -> Bool
hasSequencedLeaf :: EpidemicTree -> Bool
hasSequencedLeaf Shoot {} = Bool
False
hasSequencedLeaf (Leaf EpidemicEvent
e) =
  case EpidemicEvent
e of
    IndividualSample {Bool
Person
AbsoluteTime
indSampSeq :: Bool
indSampPerson :: Person
indSampTime :: AbsoluteTime
indSampSeq :: EpidemicEvent -> Bool
indSampPerson :: EpidemicEvent -> Person
indSampTime :: EpidemicEvent -> AbsoluteTime
..} -> Bool
indSampSeq
    PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
..} -> Bool
popSampSeq
    EpidemicEvent
_ -> Bool
False
hasSequencedLeaf (Branch EpidemicEvent
_ EpidemicTree
lt EpidemicTree
rt) = EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
lt Bool -> Bool -> Bool
|| EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
rt

-- | The events that were observed during the epidemic, ie those in the
-- reconstructed tree and any unsequenced samples. If this is not possible an
-- error message will be returned.
observedEvents :: [EpidemicEvent] -> Either String [Observation]
observedEvents :: [EpidemicEvent] -> Either String [Observation]
observedEvents [EpidemicEvent]
epiEvents = do
  EpidemicTree
epiTree <- [EpidemicEvent] -> Either String EpidemicTree
maybeEpidemicTree [EpidemicEvent]
epiEvents
  let (PointProcessEvents [Observation]
unseqObss) = EpidemicTree -> PointProcessEvents
pointProcessEvents EpidemicTree
epiTree
  [Observation]
reconTreeEvents <-
    if EpidemicTree -> Bool
hasSequencedLeaf EpidemicTree
epiTree
      then ((ReconstructedTree -> [Observation])
-> Either String ReconstructedTree -> Either String [Observation]
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM ReconstructedTree -> [Observation]
reconstructedTreeEvents) (Either String ReconstructedTree -> Either String [Observation])
-> Either String ReconstructedTree -> Either String [Observation]
forall a b. (a -> b) -> a -> b
$ EpidemicTree -> Either String ReconstructedTree
maybeReconstructedTree EpidemicTree
epiTree
      else [Observation] -> Either String [Observation]
forall a b. b -> Either a b
Right []
  [Observation] -> Either String [Observation]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Observation] -> Either String [Observation])
-> [Observation] -> Either String [Observation]
forall a b. (a -> b) -> a -> b
$ [Observation] -> [Observation]
forall a. Ord a => [a] -> [a]
List.sort ([Observation] -> [Observation])
-> ([Observation] -> [Observation])
-> [Observation]
-> [Observation]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Observation] -> [Observation]
forall a. Eq a => [a] -> [a]
List.nub ([Observation] -> [Observation]) -> [Observation] -> [Observation]
forall a b. (a -> b) -> a -> b
$ [Observation]
unseqObss [Observation] -> [Observation] -> [Observation]
forall a. [a] -> [a] -> [a]
++ [Observation]
reconTreeEvents

-- | A sorted list of all of the observations in the reconstructed tree.
reconstructedTreeEvents :: ReconstructedTree -> [Observation]
reconstructedTreeEvents :: ReconstructedTree -> [Observation]
reconstructedTreeEvents ReconstructedTree
rt =
  case ReconstructedTree
rt of
    RBranch Observation
obs ReconstructedTree
rtl ReconstructedTree
rtr ->
      [Observation] -> [Observation]
forall a. Ord a => [a] -> [a]
List.sort ([Observation] -> [Observation]) -> [Observation] -> [Observation]
forall a b. (a -> b) -> a -> b
$
      Observation
obs Observation -> [Observation] -> [Observation]
forall a. a -> [a] -> [a]
: (ReconstructedTree -> [Observation]
reconstructedTreeEvents ReconstructedTree
rtl [Observation] -> [Observation] -> [Observation]
forall a. [a] -> [a] -> [a]
++ ReconstructedTree -> [Observation]
reconstructedTreeEvents ReconstructedTree
rtr)
    RLeaf Observation
obs -> [Observation
obs]