packages feed

epi-sim-0.7.0: test/Spec.hs

{-# LANGUAGE OverloadedStrings #-}

import           Control.Exception                  (evaluate)
import           Control.Monad
import qualified Data.Aeson                         as Json
import qualified Data.ByteString                    as B
import qualified Data.ByteString.Builder            as BBuilder
import           Data.Either                        (isRight)
import           Data.Maybe                         (fromJust, isJust,
                                                     isNothing)
import qualified Data.Vector                        as V
import           Epidemic
import qualified Epidemic.Model.BDSCOD              as BDSCOD
import qualified Epidemic.Model.InhomogeneousBDS    as InhomBDS
import qualified Epidemic.Model.InhomogeneousBDSCOD as InhomBDSCOD
import           Epidemic.Types.Events
import           Epidemic.Types.Newick
import           Epidemic.Types.Observations
import           Epidemic.Types.Parameter
import           Epidemic.Types.Population
import           Epidemic.Types.Simulation          (SimulationState (..),
                                                     TerminationHandler (..),
                                                     genIOFromFixed,
                                                     genIOFromSystem,
                                                     genIOFromWord32)
import           Epidemic.Types.Time
import           Epidemic.Utility
import           Statistics.Sample
import qualified System.Random.MWC                  as MWC
import           Test.Hspec

-- | Helper function for converting from Either to Maybe monad.
either2Maybe x = case x of
  Right v -> Just v
  Left _  -> Nothing

-- | y is within n% of x from x.
withinNPercent n x y = x - d < y && y < x + d
  where
    d = n * x / 100

p1 = Person (Identifier 1)

p2 = Person (Identifier 2)

p3 = Person (Identifier 3)

p4 = Person (Identifier 4)

p5 = Person (Identifier 5)

p6 = Person (Identifier 6)

p7 = Person (Identifier 7)

-- | The first set of test data does not have any catastrophe events.
demoFullEvents01 =
  [ Infection (AbsoluteTime 1) p1 p2
  , Infection (AbsoluteTime 2) p1 p3
  , IndividualSample (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , Infection (AbsoluteTime 5) p2 p5
  , IndividualSample (AbsoluteTime 6) p4 True
  , Infection (AbsoluteTime 7) p3 p6
  , IndividualSample (AbsoluteTime 8) p2 False
  , Removal (AbsoluteTime 9) p3
  , Infection (AbsoluteTime 10) p5 p7
  , IndividualSample (AbsoluteTime 11) p6 False
  , IndividualSample (AbsoluteTime 12) p5 True
  , Removal (AbsoluteTime 13) p7
  ]

demoSampleEvents01 =
  [ Infection (AbsoluteTime 1) p1 p2
  , IndividualSample (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , IndividualSample (AbsoluteTime 6) p4 True
  , IndividualSample (AbsoluteTime 8) p2 False
  , IndividualSample (AbsoluteTime 11) p6 False
  , IndividualSample (AbsoluteTime 12) p5 True
  ]

-- | The second set of test data is the same as the first but includes a
-- catastrophe event.
demoFullEvents02 =
  [ Infection (AbsoluteTime 1) p1 p2
  , Infection (AbsoluteTime 2) p1 p3
  , IndividualSample (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , Infection (AbsoluteTime 5) p2 p5
  , IndividualSample (AbsoluteTime 6) p4 True
  , Infection (AbsoluteTime 7) p3 p6
  , IndividualSample (AbsoluteTime 8) p2 False
  , Removal (AbsoluteTime 9) p3
  , Infection (AbsoluteTime 10) p5 p7
  , PopulationSample (AbsoluteTime 11) (asPeople [p5]) True
  , IndividualSample (AbsoluteTime 12) p6 False
  , Removal (AbsoluteTime 13) p7
  ]

demoSampleEvents02 =
  [ Infection (AbsoluteTime 1) p1 p2
  , IndividualSample (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , IndividualSample (AbsoluteTime 6) p4 True
  , IndividualSample (AbsoluteTime 8) p2 False
  , PopulationSample (AbsoluteTime 11) (asPeople [p5]) True
  , IndividualSample (AbsoluteTime 12) p6 False
  ]

-- | Another test set to test that catastrophes are handled correctly.
demoFullEvents03 =
  [ Infection (AbsoluteTime 1) p1 p4
  , Infection (AbsoluteTime 2) p1 p2
  , IndividualSample (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p3
  , Infection (AbsoluteTime 5) p4 p5
  , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True
  ]

demoSampleEvents03 =
  [ Infection (AbsoluteTime 1) p1 p4
  , Infection (AbsoluteTime 2) p1 p2
  , IndividualSample (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p3
  , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True
  ]

-- | Another test to make sure that disasters are handled.
demoFullEvents04 =
  [ Infection (AbsoluteTime 1) p1 p4
  , Infection (AbsoluteTime 2) p1 p2
  , IndividualSample (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p3
  , Infection (AbsoluteTime 5) p4 p5
  , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True
  , Infection (AbsoluteTime 7) p5 p6
  , Infection (AbsoluteTime 8) p5 p7
  , PopulationSample (AbsoluteTime 9) (asPeople [p5, p6]) False
  ]

demoSampleEvents04 =
  [ Infection (AbsoluteTime 1) p1 p4
  , Infection (AbsoluteTime 2) p1 p2
  , IndividualSample (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p3
  , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True
  , PopulationSample (AbsoluteTime 9) (asPeople [p5, p6]) False
  ]

eventHandlingTests = do
  describe "Post-simulation processing" $ do
    it "Extracting observed events" $ do
      let demoEvents =
            [ PopulationSample (AbsoluteTime 0.5) (asPeople []) True -- Because the first event is a null event it can be ignored!
            , Infection (AbsoluteTime 1.0) p1 p2
            , PopulationSample (AbsoluteTime 1.5) (asPeople []) True
            , PopulationSample (AbsoluteTime 2.0) (asPeople [p1, p2]) True
            ]
      (length demoEvents == 4) `shouldBe` True
      ((length <$> observedEvents (tail demoEvents)) == (Right 2)) `shouldBe`
        True
      ((length <$> observedEvents (demoEvents)) == (Right 2)) `shouldBe` True
      (observedEvents (demoEvents) == observedEvents (tail demoEvents)) `shouldBe`
        True
      (maybeEpidemicTree (demoEvents) == maybeEpidemicTree (tail demoEvents)) `shouldBe`
        True
  describe "Catastrophe definitions" $ do
    it "Check we can find a catastrophe" $ do
      (noScheduledEvent (AbsoluteTime 0) (AbsoluteTime 1) (Timed [])) `shouldBe`
        True
      (noScheduledEvent
         (AbsoluteTime 0)
         (AbsoluteTime 1)
         (Timed [(AbsoluteTime 2, 0.5)])) `shouldBe`
        True
      (noScheduledEvent
         (AbsoluteTime 0)
         (AbsoluteTime 1)
         (Timed [(AbsoluteTime 0.5, 0.5)])) `shouldBe`
        False
      (noScheduledEvent
         (AbsoluteTime 0)
         (AbsoluteTime 1)
         (Timed [(AbsoluteTime 2, 0.6), (AbsoluteTime 0.5, 0.5)])) `shouldBe`
        False
    it "Check we can find a particular catastrophe" $ do
      (firstScheduled (AbsoluteTime 1) (Timed [])) `shouldBe` Nothing
      (firstScheduled (AbsoluteTime 1) (Timed [(AbsoluteTime 2, 0.5)])) `shouldBe`
        Just (AbsoluteTime 2, 0.5)
      (firstScheduled (AbsoluteTime 1) (Timed [(AbsoluteTime 0.5, 0.5)])) `shouldBe`
        Nothing
      (firstScheduled
         (AbsoluteTime 1)
         (Timed [(AbsoluteTime 2, 0.6), (AbsoluteTime 0.5, 0.5)])) `shouldBe`
        Just (AbsoluteTime 2, 0.6)
      isNothing
        (asTimed
           [ (AbsoluteTime 2, 0.6 :: Rate)
           , (AbsoluteTime 0.5, 0.5)
           , (AbsoluteTime 1.5, 0.4)
           ]) `shouldBe`
        True
      (firstScheduled
         (AbsoluteTime 1)
         (Timed
            [ (AbsoluteTime 2, 0.6)
            , (AbsoluteTime 0.5, 0.5)
            , (AbsoluteTime 1.5, 0.4)
            ])) `shouldBe`
        Just (AbsoluteTime 2, 0.6)
    it "Works on a very specific case it seems to not like" $ do
      (noScheduledEvent
         (AbsoluteTime 2.28)
         (AbsoluteTime (2.28 + 0.42))
         (Timed [(AbsoluteTime 2.3, 0.9)])) `shouldBe`
        False
  describe "Disaster definitions" $ do
    it "Disasters are handled correctly" $ do
      ((Right $ [Observation e | e <- demoSampleEvents04]) ==
       (observedEvents demoFullEvents04)) `shouldBe`
        True
    it "Disasters can be simulated" $ do
      (Right demoSim) <-
        simulationWithFixedSeed
          (fromJust
             (BDSCOD.configuration
                (TimeDelta 4)
                False
                Nothing
                ( 1.3
                , 0.1
                , 0.1
                , [(AbsoluteTime 3, 0.5)]
                , 0.2
                , [(AbsoluteTime 3.5, 0.5)])))
          (allEvents BDSCOD.randomEvent)
      length demoSim > 1 `shouldBe` True
  describe "Extracting observed events" $ do
    it "unseq obs still extracted if no seq obs" $ do
      let noSequencedEvents =
            [ Infection (AbsoluteTime 4.1) p1 p2
            , Infection (AbsoluteTime 4.3) p2 p3
            , Infection (AbsoluteTime 4.5) p2 p4
            , IndividualSample
                { indSampTime = AbsoluteTime 5.3
                , indSampPerson = p1
                , indSampSeq = False
                }
            , StoppingTime (AbsoluteTime 6.0)
            ]
          expectedObs =
            [ Observation
                (IndividualSample
                   { indSampTime = AbsoluteTime 5.3
                   , indSampPerson = p1
                   , indSampSeq = False
                   })
            ]
      isRight (observedEvents noSequencedEvents) `shouldBe` True
      ((Right expectedObs) == (observedEvents noSequencedEvents)) `shouldBe`
        True

helperFuncTests = do
  describe "Helpers in Utility" $ do
    it "the isAscending function works" $ do
      (isAscending ([] :: [AbsoluteTime])) `shouldBe` True
      (isAscending [-1.0]) `shouldBe` True
      (isAscending [1.0]) `shouldBe` True
      (isAscending [1.0, 2.0]) `shouldBe` True
      (isAscending [1.0, 2.0, 3.0]) `shouldBe` True
      (isAscending [1.0, -2.0]) `shouldBe` False
      (isAscending [1.0, -2.0, 3.0]) `shouldBe` False
      (isAscending [1.0, 2.0, -3.0]) `shouldBe` False
    it "the asTimed function works" $ do
      (isJust $ asTimed []) `shouldBe` True
      (isJust $ asTimed [(AbsoluteTime 0, 1)]) `shouldBe` True
      (isJust $ asTimed [(AbsoluteTime 0, 1), (AbsoluteTime 1, 3)]) `shouldBe`
        True
      (isJust $ asTimed [(AbsoluteTime 0, 3), (AbsoluteTime 1, 1)]) `shouldBe`
        True
      (isJust $ asTimed [(AbsoluteTime 1, 3), (AbsoluteTime 0, 1)]) `shouldBe`
        False
    let demoTimed =
          fromJust $
          asTimed
            [ (AbsoluteTime 0, 1.2)
            , (AbsoluteTime 1, 3.1)
            , (AbsoluteTime 2, 2.7)
            ]
     in do it "the cadlagValue function works" $ do
             (isJust $ cadlagValue demoTimed (AbsoluteTime (-1.0))) `shouldBe`
               False
             ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.0)) `shouldBe`
               True
             ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.5)) `shouldBe`
               True
             ((== 3.1) . fromJust $ cadlagValue demoTimed (AbsoluteTime 1.5)) `shouldBe`
               True
           it "the diracDeltaValue function works" $ do
             ((== 1.2) . fromJust $ diracDeltaValue demoTimed (AbsoluteTime 0)) `shouldBe`
               True
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1)) `shouldBe`
               True
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 0.9)) `shouldBe`
               False
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1.1)) `shouldBe`
               False
           it "the hasTime function works" $ do
             (hasTime demoTimed (AbsoluteTime 0)) `shouldBe` True
             (hasTime demoTimed (AbsoluteTime 0.5)) `shouldBe` False
             (hasTime demoTimed (AbsoluteTime 1)) `shouldBe` True
             (hasTime demoTimed (AbsoluteTime 1.5)) `shouldBe` False
           it "the nextTime function works" $ do
             (AbsoluteTime 0 ==
              (fromJust $ nextTime demoTimed (AbsoluteTime (-1)))) `shouldBe`
               True
             (AbsoluteTime 1 == (fromJust $ nextTime demoTimed (AbsoluteTime 0))) `shouldBe`
               True
             (AbsoluteTime 1 ==
              (fromJust $ nextTime demoTimed (AbsoluteTime 0.5))) `shouldBe`
               True
           it "the nextTime function handles the last time correctly" $ do
             isJust (nextTime demoTimed (AbsoluteTime 1.9)) `shouldBe` True
             isJust (nextTime demoTimed (AbsoluteTime 2.0)) `shouldBe` True
             isJust (nextTime demoTimed (AbsoluteTime 2.1)) `shouldBe` True
             isJust (nextTime demoTimed (AbsoluteTime 10.0)) `shouldBe` True
           it "the maybeNextTimed function works as expected" $ do
             let (Just tA) = asTimed [(AbsoluteTime 1, (1.1 :: Double)), (AbsoluteTime 3, 2.3)]
                 (Just tB) = asTimed [(AbsoluteTime 2, (1 :: Int))]
             maybeNextTimed tA tB (AbsoluteTime 0.5) == Just (AbsoluteTime 1.0,Left 1.1) `shouldBe` True
             maybeNextTimed tA tB (AbsoluteTime 1.5) == Just (AbsoluteTime 2.0,Right 1) `shouldBe` True
             maybeNextTimed tA tB (AbsoluteTime 2.5) == Just (AbsoluteTime 3.0,Left 2.3) `shouldBe` True
             isNothing (maybeNextTimed tA tB (AbsoluteTime 3.5)) `shouldBe` True

    it "shifted times work" $
      let sf =
            fromJust $
            asTimed [(AbsoluteTime (-1.0), 2.0), (AbsoluteTime 1, 3.0)]
          val1 = cadlagValue sf (AbsoluteTime 0)
          val2 = cadlagValue sf (AbsoluteTime (-2.0))
          val3 = cadlagValue sf (AbsoluteTime 1.5)
       in do isJust val1 `shouldBe` True
             val1 == Just 2.0 `shouldBe` True
             (not $ isJust val2) `shouldBe` True
             isJust val3 `shouldBe` True
             val3 == Just 3.0 `shouldBe` True
    it "the asTimed function returns nothing as expected" $ do
      (isJust $ asTimed [(AbsoluteTime 0.0, -1)]) `shouldBe` True
      (isJust $ asTimed [(AbsoluteTime 0.0, 1), (AbsoluteTime 1.0, -1)]) `shouldBe`
        True
      let (Just timedBirthRate) =
            asTimed [(AbsoluteTime 0.0, 1.0), (AbsoluteTime 1.0, -1.0)]
      (isJust $ InhomBDS.inhomBDSRates timedBirthRate 0.5 0.5) `shouldBe` False

simTypeTests =
  describe "Test Types.Simulation PRNG helpers" $
  do it "check genIOFromFixed always gives same result" $
       do g1 <- genIOFromFixed
          u11 <- MWC.uniform g1 :: IO Double
          u12 <- MWC.uniform g1 :: IO Double

          g2 <- genIOFromFixed
          u21 <- MWC.uniform g2 :: IO Double
          u22 <- MWC.uniform g2 :: IO Double

          u11 == u21 `shouldBe` True
          u11 /= u22 `shouldBe` True
          u12 == u22 `shouldBe` True

     it "check genIOFromSystem always gives different results" $
       do g1 <- genIOFromSystem
          u11 <- MWC.uniform g1 :: IO Double
          u12 <- MWC.uniform g1 :: IO Double

          g2 <- genIOFromSystem
          u21 <- MWC.uniform g2 :: IO Double
          u22 <- MWC.uniform g2 :: IO Double

          u11 /= u12 `shouldBe` True
          u11 /= u21 `shouldBe` True
          u11 /= u22 `shouldBe` True
          u12 /= u21 `shouldBe` True
          u12 /= u22 `shouldBe` True
          u21 /= u22 `shouldBe` True

     it "check genIOFromWord32 works as expected" $
       do g1 <- genIOFromWord32 1
          u11 <- MWC.uniform g1 :: IO Double
          u12 <- MWC.uniform g1 :: IO Double

          g2 <- genIOFromWord32 1
          u21 <- MWC.uniform g2 :: IO Double
          u22 <- MWC.uniform g2 :: IO Double

          g3 <- genIOFromWord32 2
          u31 <- MWC.uniform g3 :: IO Double
          u32 <- MWC.uniform g3 :: IO Double

          u11 == u21 `shouldBe` True
          u11 /= u22 `shouldBe` True
          u12 == u22 `shouldBe` True

          u11 /= u12 `shouldBe` True
          u11 /= u31 `shouldBe` True
          u11 /= u32 `shouldBe` True
          u12 /= u31 `shouldBe` True
          u12 /= u32 `shouldBe` True
          u31 /= u32 `shouldBe` True

inhomExpTests =
  describe "Test the inhomogeneous exponential variate generator" $
  let rate1 = 2.0
      sF1 = fromJust $ asTimed [(AbsoluteTime 0, rate1)]
      mean1 = 1 / rate1
      var1 = 1 / (rate1 ** 2.0)
      sF2 =
        fromJust $ asTimed [(AbsoluteTime 0, 1e-10), (AbsoluteTime 1, rate1)]
      mean2 = 1 / rate1 + 1
      var2 = var1
   in do it "check we can get a positive variate out" $ do
           gen <- genIOFromSystem
           u1 <- MWC.uniform gen :: IO Double
           (u1 > 0) `shouldBe` True
           (Just x1) <- inhomExponential sF1 (AbsoluteTime 0) gen
           (x1 > AbsoluteTime 0) `shouldBe` True
           (x1 < AbsoluteTime 100) `shouldBe` True
           True `shouldBe` True
         it "check the mean and variance look sensible" $ do
           gen <- genIOFromSystem
           xBoxed <-
             V.replicateM 20000 (inhomExponential sF1 (AbsoluteTime 0) gen)
           let x = fmap (\(Just (AbsoluteTime t)) -> t) xBoxed
           withinNPercent 5 (mean x) mean1 `shouldBe` True
           withinNPercent 5 (variance x) var1 `shouldBe` True
         it "check the mean and variance look sensible with delay" $ do
           gen <- genIOFromSystem
           xBoxed <-
             V.replicateM 20000 (inhomExponential sF2 (AbsoluteTime 0) gen)
           let x = fmap (\(Just (AbsoluteTime t)) -> t) xBoxed
           withinNPercent 5 (mean x) mean2 `shouldBe` True
           withinNPercent 5 (variance x) var2 `shouldBe` True

illFormedTreeTest :: SpecWith ()
illFormedTreeTest =
  describe "Prevent the simulator returning a broken tree" $ do
    let simDuration = TimeDelta 0.2
        simLambda = 3.2
        simMu = 0.3
        simPsi = 0.3
        simRho = 0.15
        simRhoTime = AbsoluteTime 2.6
        simOmega = 0.3
        simNu = 0.15
        simNuTime = AbsoluteTime 3.0
        simParams =
          ( simLambda
          , simMu
          , simPsi
          , [(simRhoTime, simRho)]
          , simOmega
          , [(simNuTime, simNu)])
        simConfig = BDSCOD.configuration simDuration True Nothing simParams
     in do it "stress testing the observed events function" $ do
             null (observedEvents []) `shouldBe` True
             (Right simEvents) <-
               simulationWithFixedSeed (fromJust simConfig) (allEvents BDSCOD.randomEvent)
             any isReconTreeLeaf simEvents `shouldBe` True
             let (Right oes) = observedEvents simEvents
             (length oes > 1) `shouldBe` True
           it "check leaves are recognised correctly" $ do
             let (absT, person) = (AbsoluteTime 1.0, Person (Identifier 1))
                 infEvent = Infection absT person person
                 remEvent = Removal absT person
                 sampIndv = IndividualSample absT person
                 popSampEmpty = PopulationSample absT $ asPeople []
                 popSampPerson = PopulationSample absT $ asPeople [person]
             isReconTreeLeaf infEvent `shouldBe` False
             isReconTreeLeaf remEvent `shouldBe` False
             isReconTreeLeaf (sampIndv True) `shouldBe` True
             isReconTreeLeaf (sampIndv False) `shouldBe` False
             isReconTreeLeaf (popSampEmpty True) `shouldBe` False
             isReconTreeLeaf (popSampEmpty False) `shouldBe` False
             isReconTreeLeaf (popSampPerson True) `shouldBe` True
             isReconTreeLeaf (popSampPerson False) `shouldBe` False


resultAA = demoSampleEvents01

resultAB =
  [ Infection (AbsoluteTime 1) p1 p2
  , IndividualSample (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , IndividualSample (AbsoluteTime 6) p4 True
  , PopulationSample (AbsoluteTime 12) (asPeople [p2, p6]) False
  , IndividualSample (AbsoluteTime 12) p5 True
  , PopulationSample (AbsoluteTime 17.0) (asPeople []) False
  ]

resultBA =
  [ Infection (AbsoluteTime 1) p1 p2
  , IndividualSample (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , IndividualSample (AbsoluteTime 6) p4 True
  , IndividualSample (AbsoluteTime 8) p2 False
  , IndividualSample (AbsoluteTime 11) p6 False
  , PopulationSample (AbsoluteTime 13) (asPeople [p5]) True
  ]

resultBB =
  [ Infection (AbsoluteTime 1.0) p1 p2
  , IndividualSample (AbsoluteTime 3.0) (Person (Identifier 1)) True
  , Infection (AbsoluteTime 4.0) p2 p4
  , IndividualSample (AbsoluteTime 6.0) (Person (Identifier 4)) True
  , PopulationSample (AbsoluteTime 12) (asPeople [p2, p6]) False
  , PopulationSample (AbsoluteTime 13) (asPeople [p5]) True
  , PopulationSample (AbsoluteTime 17) (asPeople []) False
  ]

aggregationTests =
  describe "Aggregation functionality tests" $ do
    it "check it does nothing unless it needs to" $
      let demoObs1 = [Observation ee | ee <- demoFullEvents01]
          demoObs2 = [Observation ee | ee <- demoSampleEvents01]
      in do
        (aggregated [] [] demoObs1) == demoObs1 `shouldBe` True
        (aggregated [] [] demoObs2) == demoObs2 `shouldBe` True
    it "check relevant intervals are processed correctly" $
      let demoObs = [Observation ee | ee <- demoSampleEvents01]
          demoSeqInts = asConsecutiveIntervals1 [AbsoluteTime 10, AbsoluteTime 13]
          demoUnseqInts = asConsecutiveIntervals1 [AbsoluteTime 7, AbsoluteTime 12, AbsoluteTime 17]
      in do
        aggregated [] [] demoObs == (map Observation resultAA) `shouldBe` True
        aggregated [] demoUnseqInts demoObs == (map Observation resultAB) `shouldBe` True
        aggregated demoSeqInts [] demoObs == (map Observation resultBA) `shouldBe` True
        aggregated demoSeqInts demoUnseqInts demoObs == (map Observation resultBB) `shouldBe` True


inhomogeneousBDSTest =
  describe "InhomogeneousBDS module tests" $ do
    it "Check the observedEvents filters out removals" $
      let demoAllEvents =
            [ Infection (AbsoluteTime 0.1) p1 p2
            , IndividualSample (AbsoluteTime 0.2) p1 True
            , Removal (AbsoluteTime 0.3) p3
            , IndividualSample (AbsoluteTime 0.4) p2 True
            ]
          demoObsEvents =
            [ Infection (AbsoluteTime 0.1) p1 p2
            , IndividualSample (AbsoluteTime 0.2) p1 True
            , IndividualSample (AbsoluteTime 0.4) p2 True
            ]
          compObsEvents = observedEvents demoAllEvents
       in do (compObsEvents == (Right [Observation e | e <- demoObsEvents])) `shouldBe` True

helperTypeTests = do
  describe "Helpers for working with the types" $ do
    it "the isAscending function works" $ do
      (isAscending ([] :: [AbsoluteTime])) `shouldBe` True
      (isAscending [-1.0]) `shouldBe` True
      (isAscending [1.0]) `shouldBe` True
      (isAscending [1.0, 2.0]) `shouldBe` True
      (isAscending [1.0, 2.0, 3.0]) `shouldBe` True
      (isAscending [1.0, -2.0]) `shouldBe` False
      (isAscending [1.0, -2.0, 3.0]) `shouldBe` False
      (isAscending [1.0, 2.0, -3.0]) `shouldBe` False
    it "the asTimed function works" $ do
      (isJust $ asTimed []) `shouldBe` True
      (isJust $ asTimed [(AbsoluteTime 0, 1)]) `shouldBe` True
      (isJust $ asTimed [(AbsoluteTime 0, 1), (AbsoluteTime 1, 3)]) `shouldBe`
        True
      (isJust $ asTimed [(AbsoluteTime 0, 3), (AbsoluteTime 1, 1)]) `shouldBe`
        True
      (isJust $ asTimed [(AbsoluteTime 1, 3), (AbsoluteTime 0, 1)]) `shouldBe`
        False
    let demoTimed =
          fromJust $
          asTimed
            [ (AbsoluteTime 0, 1.2)
            , (AbsoluteTime 1, 3.1)
            , (AbsoluteTime 2, 2.7)
            ]
     in do it "the cadlagValue function works" $ do
             (isJust $ cadlagValue demoTimed (AbsoluteTime (-1.0))) `shouldBe`
               False
             ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.0)) `shouldBe`
               True
             ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.5)) `shouldBe`
               True
             ((== 3.1) . fromJust $ cadlagValue demoTimed (AbsoluteTime 1.5)) `shouldBe`
               True
           it "the diracDeltaValue function works" $ do
             ((== 1.2) . fromJust $ diracDeltaValue demoTimed (AbsoluteTime 0)) `shouldBe`
               True
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1)) `shouldBe`
               True
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 0.9)) `shouldBe`
               False
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1.1)) `shouldBe`
               False
           it "the hasTime function works" $ do
             (hasTime demoTimed (AbsoluteTime 0)) `shouldBe` True
             (hasTime demoTimed (AbsoluteTime 0.5)) `shouldBe` False
             (hasTime demoTimed (AbsoluteTime 1)) `shouldBe` True
             (hasTime demoTimed (AbsoluteTime 1.5)) `shouldBe` False
           it "the nextTime function works" $ do
             (AbsoluteTime 0 ==
              (fromJust $ nextTime demoTimed (AbsoluteTime (-1)))) `shouldBe`
               True
             (AbsoluteTime 1 == (fromJust $ nextTime demoTimed (AbsoluteTime 0))) `shouldBe`
               True
             (AbsoluteTime 1 ==
              (fromJust $ nextTime demoTimed (AbsoluteTime 0.5))) `shouldBe`
               True
    it "shifted times work" $
      let sf =
            fromJust $
            asTimed [(AbsoluteTime (-1.0), 2.0), (AbsoluteTime 1, 3.0)]
          val1 = cadlagValue sf (AbsoluteTime 0)
          val2 = cadlagValue sf (AbsoluteTime (-2.0))
          val3 = cadlagValue sf (AbsoluteTime 1.5)
       in do isJust val1 `shouldBe` True
             val1 == Just 2.0 `shouldBe` True
             (not $ isJust val2) `shouldBe` True
             isJust val3 `shouldBe` True
             val3 == Just 3.0 `shouldBe` True

jsonTests = do
  describe "Converting to and from JSON" $ do
    it "Conversion of Timed Rate" $ do
      let demoObj =
            Timed [(AbsoluteTime 0.0, 1.0), (AbsoluteTime 1.0, 1.0)] :: Timed Rate
          (Timed demoVals) = demoObj
          demoJson = "[[0,1],[1,1]]"
          encodedObj = Json.encode demoObj
          decodedJson = Json.decode demoJson :: Maybe (Timed Rate)
       in do True `shouldBe` True
             let (Timed foo) = demoObj
              in demoVals == foo `shouldBe` True
             encodedObj == demoJson `shouldBe` True
             isJust decodedJson `shouldBe` True
             let (Timed bar) = fromJust decodedJson
              in demoVals == bar `shouldBe` True

equalBuilders :: BBuilder.Builder -> BBuilder.Builder -> Bool
equalBuilders a b = BBuilder.toLazyByteString a == BBuilder.toLazyByteString b

newickTests =
  let p1 = Person (Identifier 1)
      p2 = Person (Identifier 2)
      p3 = Person (Identifier 3)
      ps = asPeople [p1, p2]
      maybeEpiTree =
        maybeEpidemicTree
          [ Infection (AbsoluteTime 1) p1 p2
          , Infection (AbsoluteTime 2) p2 p3
          , PopulationSample (AbsoluteTime 3) (asPeople [p1, p3]) True
          , Removal (AbsoluteTime 4) p2
          ]
      maybeEpiTree' =
        maybeEpidemicTree
          [ Infection (AbsoluteTime 1) p1 p2
          , Infection (AbsoluteTime 2) p2 p3
          , PopulationSample (AbsoluteTime 3) (asPeople [p1, p3]) True
          , IndividualSample (AbsoluteTime 4) p2 True
          ]
      maybeEpiTree'' =
        maybeEpidemicTree
          [ Infection (AbsoluteTime 1) p1 p2
          , Infection (AbsoluteTime 2) p2 p3
          , PopulationSample (AbsoluteTime 3) (asPeople [p1, p3]) False
          , IndividualSample (AbsoluteTime 4) p2 True
          ]
   in describe "Writing to Newick" $ do
        it "equalBuilders works as expected" $ do
          equalBuilders (BBuilder.charUtf8 ':') (BBuilder.charUtf8 ':') `shouldBe`
            True
          equalBuilders (BBuilder.charUtf8 'a') (BBuilder.charUtf8 ':') `shouldBe`
            False
        it "derivedFrom works as expected" $ do
          let p1 = Person (Identifier 1)
          let p2 = Person (Identifier 2)
          let p3 = Person (Identifier 3)
          let e = [Infection (AbsoluteTime 0.3) p1 p2]
          derivedFrom p1 e == derivedFrom p2 e `shouldBe` True
          derivedFrom p1 e /= derivedFrom p3 e `shouldBe` True
          derivedFrom p1 e /= [] `shouldBe` True
          null (derivedFrom p3 e) `shouldBe` True
          derivedFrom p1 e == e `shouldBe` True
          let foo =
                derivedFrom
                  (Person (Identifier 1))
                  [ Infection
                      (AbsoluteTime 0.3)
                      (Person (Identifier 1))
                      (Person (Identifier 2))
                  , IndividualSample
                      (AbsoluteTime 0.7)
                      (Person (Identifier 1))
                      True
                  ]
          let bar =
                derivedFrom
                  (Person (Identifier 2))
                  [ Infection
                      (AbsoluteTime 0.3)
                      (Person (Identifier 1))
                      (Person (Identifier 2))
                  , IndividualSample
                      (AbsoluteTime 0.7)
                      (Person (Identifier 1))
                      True
                  ]
          foo == bar `shouldBe` True
        it "maybeEpidemicTree works as expected: 1" $ do
          let e1 = Removal (AbsoluteTime 1) (Person (Identifier 1))
          maybeEpidemicTree [e1] == Right (Leaf e1) `shouldBe` True
          let t1 =
                maybeEpidemicTree
                  [ Infection
                      (AbsoluteTime 0.3)
                      (Person (Identifier 1))
                      (Person (Identifier 2))
                  , IndividualSample
                      (AbsoluteTime 0.6)
                      (Person (Identifier 2))
                      True
                  , IndividualSample
                      (AbsoluteTime 0.7)
                      (Person (Identifier 1))
                      True
                  ]
          let t2 =
                Right
                  (Branch
                     (Infection
                        (AbsoluteTime 0.3)
                        (Person (Identifier 1))
                        (Person (Identifier 2)))
                     (Leaf
                        (IndividualSample
                           (AbsoluteTime 0.7)
                           (Person (Identifier 1))
                           True))
                     (Leaf
                        (IndividualSample
                           (AbsoluteTime 0.6)
                           (Person (Identifier 2))
                           True)))
          t1 == t2 `shouldBe` True
          maybeEpidemicTree
            [ Infection
                (AbsoluteTime 0.3)
                (Person (Identifier 1))
                (Person (Identifier 2))
            ] ==
            Right
              (Branch
                 (Infection
                    (AbsoluteTime 0.3)
                    (Person (Identifier 1))
                    (Person (Identifier 2)))
                 (Shoot (Person (Identifier 1)))
                 (Shoot (Person (Identifier 2)))) `shouldBe`
            True
          maybeEpidemicTree
            [ Infection
                (AbsoluteTime 0.3)
                (Person (Identifier 1))
                (Person (Identifier 2))
            , IndividualSample (AbsoluteTime 0.7) (Person (Identifier 1)) True
            ] ==
            Right
              (Branch
                 (Infection
                    (AbsoluteTime 0.3)
                    (Person (Identifier 1))
                    (Person (Identifier 2)))
                 (Leaf
                    (IndividualSample
                       (AbsoluteTime 0.7)
                       (Person (Identifier 1))
                       True))
                 (Shoot (Person (Identifier 2)))) `shouldBe`
            True
          let trickyEvents =
                [ Infection
                    (AbsoluteTime 0.3)
                    (Person (Identifier 1))
                    (Person (Identifier 2))
                , Infection
                    (AbsoluteTime 0.4)
                    (Person (Identifier 2))
                    (Person (Identifier 3))
                , IndividualSample
                    (AbsoluteTime 0.6)
                    (Person (Identifier 3))
                    True
                , IndividualSample
                    (AbsoluteTime 0.7)
                    (Person (Identifier 1))
                    True
                ]
          isRight (maybeEpidemicTree trickyEvents) `shouldBe` True
        it "maybeEpidemicTree works as expected: 2" $ do
          let p1 = Person (Identifier 1)
              p2 = Person (Identifier 2)
              demoEvents =
                [ PopulationSample (AbsoluteTime 0.5) (asPeople []) True -- Because the first event is a null event it can be ignored!
                , Infection (AbsoluteTime 1.0) p1 p2
                , PopulationSample (AbsoluteTime 1.5) (asPeople []) True
                , PopulationSample (AbsoluteTime 2.0) (asPeople [p1, p2]) True
                ]
          (length demoEvents == 4) `shouldBe` True
          (maybeEpidemicTree demoEvents == maybeEpidemicTree (tail demoEvents)) `shouldBe`
            True
        it "asNewickString works for ReconstructedTree" $ do
          isJust
            (asNewickString
               (AbsoluteTime 0, Person (Identifier 1))
               (RLeaf
                  (Observation (IndividualSample
                     (AbsoluteTime 1)
                     (Person (Identifier 1))
                     True)))) `shouldBe`
            True
          let trickyEvents =
                [ Infection
                    (AbsoluteTime 0.3)
                    (Person (Identifier 1))
                    (Person (Identifier 2))
                , Infection
                    (AbsoluteTime 0.4)
                    (Person (Identifier 2))
                    (Person (Identifier 3))
                , IndividualSample
                    (AbsoluteTime 0.6)
                    (Person (Identifier 3))
                    True
                , IndividualSample
                    (AbsoluteTime 0.7)
                    (Person (Identifier 1))
                    True
                ]
          let et = maybeEpidemicTree trickyEvents :: Either String EpidemicTree
              rt = maybeReconstructedTree =<< et :: Either String ReconstructedTree
              maybeNewickPair = asNewickString (AbsoluteTime 0, Person (Identifier 1)) =<< (either2Maybe rt)
          let newickTarget =
                BBuilder.stringUtf8 "(1:0.39999999999999997,3:0.3):0.3"
          isJust maybeNewickPair `shouldBe` True
          [ IndividualSample (AbsoluteTime 0.6) (Person (Identifier 3)) True, IndividualSample (AbsoluteTime 0.7) (Person (Identifier 1)) True] == snd (fromJust maybeNewickPair) `shouldBe` True
          equalBuilders newickTarget (fst $ fromJust maybeNewickPair) `shouldBe`
            True
          let catasNewick =
                (asNewickString
                   (AbsoluteTime 0, Person (Identifier 1))
                   (RLeaf
                      (Observation (PopulationSample
                         (AbsoluteTime 1)
                         (asPeople
                            [Person (Identifier 1), Person (Identifier 2)])
                         True))))
          let catasTarget = BBuilder.stringUtf8 "1&2:1.0"
          equalBuilders catasTarget (fst $ fromJust catasNewick) `shouldBe` True

main :: IO ()
main =
  hspec $ do
    eventHandlingTests
    helperFuncTests
    inhomExpTests
    illFormedTreeTest
    inhomogeneousBDSTest
    helperTypeTests
    jsonTests
    newickTests
    aggregationTests
    simTypeTests
    terminationTests1

terminationTests1 =
  describe "Termination handling tests: InhomogeneousBDSCOD" $ do
    let duration = TimeDelta 2.0
        birthRateSpec = [(AbsoluteTime 0.0, 1.5), (AbsoluteTime 0.5, 0.5)]
        deathRateSpec = [(AbsoluteTime 0.0, 0.4)]
        sampRateSpec = [(AbsoluteTime 0.0, 0.1)]
        occRateSpec = [(AbsoluteTime 0.0, 0.1)]
        seqSched = [(AbsoluteTime 0.9, 0.1)]
        unseqSched = [(AbsoluteTime 0.5, 0.4), (AbsoluteTime 0.75, 0.5)]
        ratesAndProbs = (birthRateSpec,deathRateSpec,sampRateSpec,seqSched,occRateSpec,unseqSched)
        conf maybeTH = fromJust $ InhomBDSCOD.configuration duration True maybeTH ratesAndProbs
        -- We need one simulation configuration for each of the termination
        -- handlers that we want to test.
        simConfigNothing  = conf Nothing
        simConfigNever = conf (Just (const False, const ()))
        simConfigAlways = conf (Just (const True, const ()))
        numDeadThreshold = 3
        simConfigSometimes = conf (Just ((>numDeadThreshold) . InhomBDSCOD.getNumRemovedByDeath,
                                         \es -> length [() | Removal _ _ <- es]))
        allEventsFunc = allEvents InhomBDSCOD.randomEvent
    it "test simulation works without hander" $
      do
        (Right esNothing) <- simulationWithFixedSeed simConfigNothing allEventsFunc
        -- There should always be at least one event in the simulation.
        (length esNothing > 0) `shouldBe` True
    it "test simulation works with handler that does not trigger" $
      do
        (Right esNothing) <- simulationWithFixedSeed simConfigNothing allEventsFunc
        (Right esNever) <- simulationWithFixedSeed simConfigNever allEventsFunc
        -- If the handler never triggers this should look the same as not having
        -- the event handler.
        (all id $ zipWith (==) esNothing esNever) `shouldBe` True
    it "test simulation works with handler that always triggers" $
      do
        -- If the handler always triggers this should always return the summary.
        replicateM_ 30
          (do esAlways <- simulationWithSystem simConfigAlways allEventsFunc
              esAlways == Left (Just ()) `shouldBe` True
           )
    it "test simulation works with handler that sometimes triggers" $
      do
        -- If the handler only sometimes triggers then we need to test both
        -- branches.
        replicateM_ 30
          (do esSometimes <- simulationWithSystem simConfigSometimes allEventsFunc
              case esSometimes of
                -- If the termination handler did not trigger the number of
                -- removals should not exceed the threshold allowed by the
                -- termination handler
                (Right es) -> do length [() | Removal _ _ <- es] <= numDeadThreshold `shouldBe` True
                -- If the termination handler did trigger then we should know
                -- exactly how many removals there was.
                (Left (Just n)) -> do n == numDeadThreshold + 1 `shouldBe` True
                (Left Nothing) -> True `shouldBe` False -- this branch should not be reached.
                )