{-# LANGUAGE RecordWildCards #-}
module Epidemic.Types.Newick where
import qualified Data.ByteString.Builder as BBuilder
import qualified Data.List as List
import qualified Data.Vector as V
import Epidemic.Types.Observations
import Epidemic.Types.Events
import Epidemic.Types.Population
import Epidemic.Types.Time
class Newick t
where
asNewickString ::
(AbsoluteTime, Person)
-> t
-> Maybe (BBuilder.Builder, [EpidemicEvent])
ampersandBuilder :: BBuilder.Builder
ampersandBuilder :: Builder
ampersandBuilder = Char -> Builder
BBuilder.charUtf8 Char
'&'
colonBuilder :: BBuilder.Builder
colonBuilder :: Builder
colonBuilder = Char -> Builder
BBuilder.charUtf8 Char
':'
leftBraceBuilder :: BBuilder.Builder
leftBraceBuilder :: Builder
leftBraceBuilder = Char -> Builder
BBuilder.charUtf8 Char
'('
rightBraceBuilder :: BBuilder.Builder
rightBraceBuilder :: Builder
rightBraceBuilder = Char -> Builder
BBuilder.charUtf8 Char
')'
commaBuilder :: BBuilder.Builder
commaBuilder :: Builder
commaBuilder = Char -> Builder
BBuilder.charUtf8 Char
','
catastrophePeopleBuilder :: People -> BBuilder.Builder
catastrophePeopleBuilder :: People -> Builder
catastrophePeopleBuilder (People Vector Person
persons) =
[Builder] -> Builder
forall a. Monoid a => [a] -> a
mconcat ([Builder] -> Builder) -> [Builder] -> Builder
forall a b. (a -> b) -> a -> b
$
Builder -> [Builder] -> [Builder]
forall a. a -> [a] -> [a]
List.intersperse Builder
ampersandBuilder [Person -> Builder
personByteString Person
p | Person
p <- Vector Person -> [Person]
forall a. Vector a -> [a]
V.toList Vector Person
persons]
instance Newick ReconstructedTree where
asNewickString :: (AbsoluteTime, Person)
-> ReconstructedTree -> Maybe (Builder, [EpidemicEvent])
asNewickString (AbsoluteTime
t, Person
_) (RLeaf (Observation EpidemicEvent
e)) =
let branchLength :: AbsoluteTime -> AbsoluteTime -> Builder
branchLength AbsoluteTime
a AbsoluteTime
b = Double -> Builder
BBuilder.doubleDec Double
td
where
(TimeDelta Double
td) = AbsoluteTime -> AbsoluteTime -> TimeDelta
timeDelta AbsoluteTime
a AbsoluteTime
b
in case EpidemicEvent
e of
IndividualSample {Bool
Person
AbsoluteTime
indSampSeq :: EpidemicEvent -> Bool
indSampPerson :: EpidemicEvent -> Person
indSampTime :: EpidemicEvent -> AbsoluteTime
indSampSeq :: Bool
indSampPerson :: Person
indSampTime :: AbsoluteTime
..} ->
if Bool
indSampSeq
then (Builder, [EpidemicEvent]) -> Maybe (Builder, [EpidemicEvent])
forall a. a -> Maybe a
Just
( Person -> Builder
personByteString Person
indSampPerson Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>
Builder
colonBuilder Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> AbsoluteTime -> AbsoluteTime -> Builder
branchLength AbsoluteTime
t AbsoluteTime
indSampTime
, [EpidemicEvent
e])
else Maybe (Builder, [EpidemicEvent])
forall a. Maybe a
Nothing
PopulationSample {Bool
People
AbsoluteTime
popSampSeq :: EpidemicEvent -> Bool
popSampPeople :: EpidemicEvent -> People
popSampTime :: EpidemicEvent -> AbsoluteTime
popSampSeq :: Bool
popSampPeople :: People
popSampTime :: AbsoluteTime
..} ->
if Bool
popSampSeq
then (Builder, [EpidemicEvent]) -> Maybe (Builder, [EpidemicEvent])
forall a. a -> Maybe a
Just
( People -> Builder
catastrophePeopleBuilder People
popSampPeople Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>
Builder
colonBuilder Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> AbsoluteTime -> AbsoluteTime -> Builder
branchLength AbsoluteTime
t AbsoluteTime
popSampTime
, [EpidemicEvent
e])
else Maybe (Builder, [EpidemicEvent])
forall a. Maybe a
Nothing
EpidemicEvent
_ -> Maybe (Builder, [EpidemicEvent])
forall a. Maybe a
Nothing
asNewickString (AbsoluteTime
t, Person
_) (RBranch (Observation EpidemicEvent
e) ReconstructedTree
lt ReconstructedTree
rt) =
case EpidemicEvent
e of
(Infection AbsoluteTime
t' Person
p1 Person
p2) -> do
(Builder
leftNS, [EpidemicEvent]
leftEs) <- (AbsoluteTime, Person)
-> ReconstructedTree -> Maybe (Builder, [EpidemicEvent])
forall t.
Newick t =>
(AbsoluteTime, Person) -> t -> Maybe (Builder, [EpidemicEvent])
asNewickString (AbsoluteTime
t', Person
p1) ReconstructedTree
lt
(Builder
rightNS, [EpidemicEvent]
rightEs) <- (AbsoluteTime, Person)
-> ReconstructedTree -> Maybe (Builder, [EpidemicEvent])
forall t.
Newick t =>
(AbsoluteTime, Person) -> t -> Maybe (Builder, [EpidemicEvent])
asNewickString (AbsoluteTime
t', Person
p2) ReconstructedTree
rt
let branchLength :: Builder
branchLength = Double -> Builder
BBuilder.doubleDec Double
td
where
(TimeDelta Double
td) = AbsoluteTime -> AbsoluteTime -> TimeDelta
timeDelta AbsoluteTime
t AbsoluteTime
t'
(Builder, [EpidemicEvent]) -> Maybe (Builder, [EpidemicEvent])
forall (m :: * -> *) a. Monad m => a -> m a
return
( Builder
leftBraceBuilder Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>
Builder
leftNS Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>
Builder
commaBuilder Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>
Builder
rightNS Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
rightBraceBuilder Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
colonBuilder Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
branchLength
, [EpidemicEvent] -> [EpidemicEvent]
forall a. Ord a => [a] -> [a]
List.sort ([EpidemicEvent] -> [EpidemicEvent])
-> [EpidemicEvent] -> [EpidemicEvent]
forall a b. (a -> b) -> a -> b
$ [EpidemicEvent]
leftEs [EpidemicEvent] -> [EpidemicEvent] -> [EpidemicEvent]
forall a. [a] -> [a] -> [a]
++ [EpidemicEvent]
rightEs)
EpidemicEvent
_ -> Maybe (Builder, [EpidemicEvent])
forall a. Maybe a
Nothing