-- XXX: Add possibility to store supplementary information about the chain.
--
-- Maybe something like Trace b; and give a function a -> b to extract
-- supplementary info.

-- XXX: Status tuned exclusively to the Metropolis-Hastings algorithm. We
-- should abstract the algorithm from the chain. For example,
--
-- @
-- data Status a b = Status { Chain a; Algorithm a b}
-- @
--
-- where a described the state space and b the auxiliary information of the
-- algorithm. This would also solve the above problem, for example in terms of
-- the Hamiltonian algorithm

-- |
-- Module      :  Mcmc.Status
-- Description :  What is an MCMC?
-- Copyright   :  (c) Dominik Schrempf 2020
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Tue May  5 18:01:15 2020.
module Mcmc.Status
  ( Status (..),
    status,
    saveWith,
    force,
    quiet,
    debug,
  )
where

import Data.Maybe
import Data.Time.Clock
import Mcmc.Item
import Mcmc.Monitor
import Mcmc.Proposal
import Mcmc.Trace
import Mcmc.Verbosity (Verbosity (..))
import Numeric.Log
import System.IO
import System.Random.MWC hiding (save)
import Prelude hiding (cycle)

-- | The 'Status' contains all information to run an MCMC chain. It is
-- constructed using the function 'status'.
data Status a = Status
  { -- MCMC related variables; saved.

    -- | The name of the MCMC chain; used as file prefix.
    name :: String,
    -- | The current 'Item' of the chain combines the current state and the
    -- current likelihood.
    item :: Item a,
    -- | The iteration is the number of completed cycles.
    iteration :: Int,
    -- | The 'Trace' of the Markov chain in reverse order, the most recent
    -- 'Item' is at the head of the list.
    trace :: Trace a,
    -- | For each 'Proposal', store the list of accepted (True) and rejected (False)
    -- proposals; for reasons of efficiency, the list is also stored in reverse
    -- order.
    acceptance :: Acceptance (Proposal a),
    -- | Number of burn in iterations; deactivate burn in with 'Nothing'.
    burnInIterations :: Maybe Int,
    -- | Auto tuning period (only during burn in); deactivate auto tuning with
    -- 'Nothing'.
    autoTuningPeriod :: Maybe Int,
    -- | Number of normal iterations excluding burn in. Note that auto tuning
    -- only happens during burn in.
    iterations :: Int,
    --
    -- Auxiliary variables; saved.

    -- | Overwrite output files? Default is 'False', change with 'force'.
    forceOverwrite :: Bool,
    -- | Save the chain with trace of given length at the end of the run?
    -- Default is no save ('Nothing'). Change with 'saveWith'.
    save :: Maybe Int,
    -- | Verbosity.
    verbosity :: Verbosity,
    -- | The random number generator.
    generator :: GenIO,
    --
    -- Auxiliary variables; not saved.

    -- | Starting time and starting iteration of chain; used to calculate
    -- run time and ETA.
    start :: Maybe (Int, UTCTime),
    -- | Handle to log file.
    logHandle :: Maybe Handle,
    --
    -- Auxiliary functions; not saved.

    -- | The prior function. The un-normalized posterior is the product of the
    -- prior and the likelihood.
    priorF :: a -> Log Double,
    -- | The likelihood function. The un-normalized posterior is the product of
    -- the prior and the likelihood.
    likelihoodF :: a -> Log Double,
    --
    -- Variables related to the algorithm; not saved.

    -- | A set of 'Proposal's form a 'Cycle'.
    cycle :: Cycle a,
    -- | A 'Monitor' observing the chain.
    monitor :: Monitor a
  }

-- | Initialize the 'Status' of a Markov chain Monte Carlo run.
status ::
  -- | Name of the Markov chain; used as file prefix.
  String ->
  -- | The prior function.
  (a -> Log Double) ->
  -- | The likelihood function.
  (a -> Log Double) ->
  -- | A list of 'Proposal's executed in forward order. The
  -- chain will be logged after each cycle.
  Cycle a ->
  -- | A 'Monitor' observing the chain.
  Monitor a ->
  -- | The initial state in the state space @a@.
  a ->
  -- | Number of burn in iterations; deactivate burn in with 'Nothing'.
  Maybe Int ->
  -- | Auto tuning period (only during burn in); deactivate
  -- auto tuning with 'Nothing'.
  Maybe Int ->
  -- | Number of normal iterations excluding burn in. Note
  -- that auto tuning only happens during burn in.
  Int ->
  -- | A source of randomness. For reproducible runs, make
  -- sure to use a generator with the same seed.
  GenIO ->
  Status a
status n p l c m x mB mT nI g
  | isJust mT && isNothing mB = error "status: Auto tuning period given, but no burn in."
  | otherwise =
    Status
      n
      i
      0
      (singletonT i)
      (emptyA $ ccProposals c)
      mB
      mT
      nI
      False
      Nothing
      Info
      g
      Nothing
      Nothing
      p
      l
      c
      m
  where
    i = Item x (p x) (l x)

-- | Save the Markov chain with trace of given length.
saveWith :: Int -> Status a -> Status a
saveWith n s = s {save = Just n}

-- | Overwrite existing files; it is not necessary to use 'force', when a chain
-- is continued.
force :: Status a -> Status a
force s = s {forceOverwrite = True}

-- | Do not print anything to standard output. Do not create log file. File
-- monitors and batch monitors are executed normally.
quiet :: Status a -> Status a
quiet s = s {verbosity = Quiet}

-- | Be verbose.
debug :: Status a -> Status a
debug s = s {verbosity = Debug}