{-# 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 of types that can be expressed in Newick format.
class Newick t
  where
  asNewickString ::
       (AbsoluteTime, Person) -- ^ The person and time of the root of the tree
    -> 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