{-# 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
      scRates           :: r
      -- | The population
    , scPopulation      :: p
      -- | A new identifier
    , scNewIdentifier   :: Identifier
      -- | The absolute time at which the simulation starts
    , scStartTime       :: AbsoluteTime
      -- | The duration of the simulation until it stops
    , scSimDuration     :: TimeDelta
      -- | The simulation terminates if this predicate is not satisfied
    , scTerminationHandler :: Maybe (TerminationHandler p s)
      -- | The simulation requires at least two sequenced samples
    , 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 (Eq, 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 seed = initialize (Unboxed.fromList [seed])

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

-- | A PRNG seed which is hard coded into @mwc-random@.
genIOFromFixed :: IO GenIO
genIOFromFixed = create