-- Note: It is not necessary to add another type @b@ to store supplementary
-- information about the chain. The information can just be stored in @a@
-- equally well.

-- XXX: Status tuned exclusively to the Metropolis-Hastings algorithm. We should
-- abstract the algorithm from the chain. Maybe something like:
--
-- @
-- data Status a = Status { Chain a; Algorithm a}
-- @

-- |
-- 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
  ( Cleaner (..),
    Status (..),
    status,
    cleanWith,
    saveWith,
    force,
    quiet,
    debug,
    noData,
  )
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)

-- | Clean the state periodically.
--
-- The prior and the likelihood will be updated after the cleaning process.
--
-- For long chains, successive numerical errors can accumulate such that the
-- state diverges from honoring specific required constraints. In these cases, a
-- 'Cleaner' can be used to ensure that the required constraints of the state
-- are honored. For example, the branches of an ultrametric phylogeny may
-- diverge slightly after successful many proposals such that the phylogeny is
-- not anymore ultrametric.
--
-- Please be aware that the Markov chain will not converge to the true posterior
-- distribution if the state is changed substantially! Only apply subtle changes
-- that are absolutely necessary to preserve the required properties of the
-- state such as specific numerical constraints.
data Cleaner a = Cleaner
  { -- | Clean every given number of iterations.
    Cleaner a -> Int
clEvery :: Int,
    -- | Cleaning function. Executed before monitoring the state.
    Cleaner a -> a -> a
clFunction :: a -> a
  }

-- | The 'Status' contains all information to run an MCMC chain. It is
-- constructed using the function 'status'.
--
-- The polymorphic type @a@ stores the state of the chain. It can also be used
-- to store auxiliary information.
data Status a = Status
  { -- MCMC related variables; saved.

    -- | The name of the MCMC chain; used as file prefix.
    Status a -> String
name :: String,
    -- | The current 'Item' of the chain combines the current state and the
    -- current likelihood.
    Status a -> Item a
item :: Item a,
    -- | The iteration is the number of completed cycles.
    Status a -> Int
iteration :: Int,
    -- | The 'Trace' of the Markov chain in reverse order, the most recent
    -- 'Item' is at the head of the list.
    Status a -> Trace a
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.
    Status a -> Acceptance (Proposal a)
acceptance :: Acceptance (Proposal a),
    -- | Number of burn in iterations; deactivate burn in with 'Nothing'.
    Status a -> Maybe Int
burnInIterations :: Maybe Int,
    -- | Auto tuning period (only during burn in); deactivate auto tuning with
    -- 'Nothing'.
    Status a -> Maybe Int
autoTuningPeriod :: Maybe Int,
    -- | Number of normal iterations excluding burn in. Note that auto tuning
    -- only happens during burn in.
    Status a -> Int
iterations :: Int,
    --
    -- Auxiliary variables; saved.

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

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

    -- | The prior function. The un-normalized posterior is the product of the
    -- prior and the likelihood.
    Status a -> a -> Log Double
priorF :: a -> Log Double,
    -- | The likelihood function. The un-normalized posterior is the product of
    -- the prior and the likelihood.
    Status a -> a -> Log Double
likelihoodF :: a -> Log Double,
    -- | Clean the state periodically.
    Status a -> Maybe (Cleaner a)
cleaner :: Maybe (Cleaner a),
    --
    -- Variables related to the algorithm; not saved.

    -- | A set of 'Proposal's form a 'Cycle'.
    Status a -> Cycle a
cycle :: Cycle a,
    -- | A 'Monitor' observing the chain.
    Status a -> Monitor a
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
  -- generators with the same, fixed seed.
  GenIO ->
  Status a
status :: String
-> (a -> Log Double)
-> (a -> Log Double)
-> Cycle a
-> Monitor a
-> a
-> Maybe Int
-> Maybe Int
-> Int
-> GenIO
-> Status a
status String
n a -> Log Double
p a -> Log Double
l Cycle a
c Monitor a
m a
x Maybe Int
mB Maybe Int
mT Int
nI GenIO
g
  | Maybe Int -> Bool
forall a. Maybe a -> Bool
isJust Maybe Int
mT Bool -> Bool -> Bool
&& Maybe Int -> Bool
forall a. Maybe a -> Bool
isNothing Maybe Int
mB = String -> Status a
forall a. HasCallStack => String -> a
error String
"status: Auto tuning period given, but no burn in."
  | Bool
otherwise =
    String
-> Item a
-> Int
-> Trace a
-> Acceptance (Proposal a)
-> Maybe Int
-> Maybe Int
-> Int
-> Bool
-> Maybe Int
-> Verbosity
-> GenIO
-> Maybe (Int, UTCTime)
-> Maybe Handle
-> (a -> Log Double)
-> (a -> Log Double)
-> Maybe (Cleaner a)
-> Cycle a
-> Monitor a
-> Status a
forall a.
String
-> Item a
-> Int
-> Trace a
-> Acceptance (Proposal a)
-> Maybe Int
-> Maybe Int
-> Int
-> Bool
-> Maybe Int
-> Verbosity
-> GenIO
-> Maybe (Int, UTCTime)
-> Maybe Handle
-> (a -> Log Double)
-> (a -> Log Double)
-> Maybe (Cleaner a)
-> Cycle a
-> Monitor a
-> Status a
Status
      String
n
      Item a
i
      Int
0
      (Item a -> Trace a
forall a. Item a -> Trace a
singletonT Item a
i)
      ([Proposal a] -> Acceptance (Proposal a)
forall k. Ord k => [k] -> Acceptance k
emptyA ([Proposal a] -> Acceptance (Proposal a))
-> [Proposal a] -> Acceptance (Proposal a)
forall a b. (a -> b) -> a -> b
$ Cycle a -> [Proposal a]
forall a. Cycle a -> [Proposal a]
ccProposals Cycle a
c)
      Maybe Int
mB
      Maybe Int
mT
      Int
nI
      Bool
False
      Maybe Int
forall a. Maybe a
Nothing
      Verbosity
Info
      GenIO
g
      Maybe (Int, UTCTime)
forall a. Maybe a
Nothing
      Maybe Handle
forall a. Maybe a
Nothing
      a -> Log Double
p
      a -> Log Double
l
      Maybe (Cleaner a)
forall a. Maybe a
Nothing
      Cycle a
c
      Monitor a
m
  where
    i :: Item a
i = a -> Log Double -> Log Double -> Item a
forall a. a -> Log Double -> Log Double -> Item a
Item a
x (a -> Log Double
p a
x) (a -> Log Double
l a
x)

-- | Clean the state every given number of generations using the given function.
-- See 'Cleaner'.
cleanWith :: Cleaner a -> Status a -> Status a
cleanWith :: Cleaner a -> Status a -> Status a
cleanWith Cleaner a
c Status a
s = Status a
s {cleaner :: Maybe (Cleaner a)
cleaner = Cleaner a -> Maybe (Cleaner a)
forall a. a -> Maybe a
Just Cleaner a
c}

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

-- | Overwrite existing files; it is not necessary to use 'force', when a chain
-- is continued.
force :: Status a -> Status a
force :: Status a -> Status a
force Status a
s = Status a
s {forceOverwrite :: Bool
forceOverwrite = Bool
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 :: Status a -> Status a
quiet Status a
s = Status a
s {verbosity :: Verbosity
verbosity = Verbosity
Quiet}

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

-- | Set the likelihood function to 1.0. Useful for debugging and testing.
noData :: Status a -> Status a
noData :: Status a -> Status a
noData Status a
s = Status a
s {likelihoodF :: a -> Log Double
likelihoodF = Log Double -> a -> Log Double
forall a b. a -> b -> a
const Log Double
1.0}