{-# LANGUAGE GADTs #-}

module Epidemic.Types.Simulation
  ( SimulationConfiguration(..)
  , SimulationState(..)
  , SimulationRandEvent(..)
  , TerminationHandler(..)
  , genIOFromFixed
  , genIOFromWord32
  , genIOFromSystem
  ) where

import qualified Data.Vector.Unboxed       as Unboxed
import           Epidemic.Types.Events
import           Epidemic.Types.Parameter
import           Epidemic.Types.Population
import           Epidemic.Types.Time       (AbsoluteTime (..), TimeDelta (..))
import           GHC.Word                  (Word32)
import           System.Random.MWC         (GenIO, create, createSystemRandom,
                                            initialize)

data SimulationConfiguration r p s =
  SimulationConfiguration
    { -- | The event rates
      SimulationConfiguration r p s -> r
scRates           :: r
      -- | The population
    , SimulationConfiguration r p s -> p
scPopulation      :: p
      -- | A new identifier
    , SimulationConfiguration r p s -> Identifier
scNewIdentifier   :: Identifier
      -- | The absolute time at which the simulation starts
    , SimulationConfiguration r p s -> AbsoluteTime
scStartTime       :: AbsoluteTime
      -- | The duration of the simulation until it stops
    , SimulationConfiguration r p s -> TimeDelta
scSimDuration     :: TimeDelta
      -- | The simulation terminates if this predicate is not satisfied
    , SimulationConfiguration r p s -> Maybe (TerminationHandler p s)
scTerminationHandler :: Maybe (TerminationHandler p s)
      -- | The simulation requires at least two sequenced samples
    , SimulationConfiguration r p s -> Bool
scRequireCherry   :: Bool
    }

-- | Either there is a valid simulation state which contains a sequence of
-- epidemic events along with the time and population or, if the simulation has
-- terminated early there is another value to indicate that along with a value
-- which can be used to indicate why the simulation was terminated early.
data SimulationState b c
  = SimulationState (AbsoluteTime, [EpidemicEvent], b, Identifier)
  | TerminatedSimulation (Maybe c)
  deriving (SimulationState b c -> SimulationState b c -> Bool
(SimulationState b c -> SimulationState b c -> Bool)
-> (SimulationState b c -> SimulationState b c -> Bool)
-> Eq (SimulationState b c)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall b c.
(Eq b, Eq c) =>
SimulationState b c -> SimulationState b c -> Bool
/= :: SimulationState b c -> SimulationState b c -> Bool
$c/= :: forall b c.
(Eq b, Eq c) =>
SimulationState b c -> SimulationState b c -> Bool
== :: SimulationState b c -> SimulationState b c -> Bool
$c== :: forall b c.
(Eq b, Eq c) =>
SimulationState b c -> SimulationState b c -> Bool
Eq, Int -> SimulationState b c -> ShowS
[SimulationState b c] -> ShowS
SimulationState b c -> String
(Int -> SimulationState b c -> ShowS)
-> (SimulationState b c -> String)
-> ([SimulationState b c] -> ShowS)
-> Show (SimulationState b c)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall b c. (Show b, Show c) => Int -> SimulationState b c -> ShowS
forall b c. (Show b, Show c) => [SimulationState b c] -> ShowS
forall b c. (Show b, Show c) => SimulationState b c -> String
showList :: [SimulationState b c] -> ShowS
$cshowList :: forall b c. (Show b, Show c) => [SimulationState b c] -> ShowS
show :: SimulationState b c -> String
$cshow :: forall b c. (Show b, Show c) => SimulationState b c -> String
showsPrec :: Int -> SimulationState b c -> ShowS
$cshowsPrec :: forall b c. (Show b, Show c) => Int -> SimulationState b c -> ShowS
Show)

data SimulationRandEvent a b where
  SimulationRandEvent
    :: (ModelParameters a b, Population b)
    => (a
    -> AbsoluteTime
    -> b
    -> Identifier
    -> GenIO
    -> IO (AbsoluteTime, EpidemicEvent, b, Identifier))
    -> SimulationRandEvent a b

-- | Check if a simulation should be terminated and if it should be terminated,
-- then compute a summary explaining why. The first function is used to
-- determine whether the population has entered a state which requires the
-- simulation to terminate early and the second can be use to write a summary of
-- the events that led to the termination.
data TerminationHandler b c where
  TerminationHandler
    :: Population b
    => (b -> Bool)
    -> ([EpidemicEvent] -> c)
    -> TerminationHandler b c

-- | A PRNG seed based on the given number. This is the best choice for
-- reproducible simulations.
genIOFromWord32 :: Word32 -> IO GenIO
genIOFromWord32 :: Word32 -> IO GenIO
genIOFromWord32 Word32
seed = Vector Word32 -> IO GenIO
forall (m :: * -> *) (v :: * -> *).
(PrimMonad m, Vector v Word32) =>
v Word32 -> m (Gen (PrimState m))
initialize ([Word32] -> Vector Word32
forall a. Unbox a => [a] -> Vector a
Unboxed.fromList [Word32
seed])

-- | A PRNG seed generated by the system's random number generator.
genIOFromSystem :: IO GenIO
genIOFromSystem :: IO GenIO
genIOFromSystem = IO GenIO
createSystemRandom

-- | A PRNG seed which is hard coded into @mwc-random@.
genIOFromFixed :: IO GenIO
genIOFromFixed :: IO GenIO
genIOFromFixed = IO GenIO
forall (m :: * -> *). PrimMonad m => m (Gen (PrimState m))
create