-- 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,
    noSave,
    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 at the end of the run? Default is 'True', change with 'noSave'.
        save :: Bool,
        -- | 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
      True
      Info
      g
      Nothing
      Nothing
      p
      l
      c
      m
  where
    i = Item x (p x) (l x)

-- | Do not save the Markov chain at the end.
noSave :: Status a -> Status a
noSave s = s {save = False}

-- | 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}