{-# LANGUAGE MultiParamTypeClasses #-}

module MCMC.Types ( 
                  -- * Targets and Proposals
                  Rand
                  , Density
                  , Sample
                  , Target
                  , viewTarget
                  , TargetView(..)
                  , makeTarget
                  , Proposal
                  , viewProposal
                  , ProposalView(..)
                  , makeProposal
                  -- * Transition kernels
                  , Step
                  , Kernel
                  -- * Actions
                  , Act
                  , Action
                  , viewAction
                  , ActionView(..)
                  , makeAction
                  ) where

import qualified System.Random.MWC as MWC

-- | An even shorter name for PRNGs in the 'IO' monad.
type Rand = MWC.GenIO

-- | The probability density function used in both target and proposal distributions.
-- Given an input point, this method returns a probability density. 
type Density a = a -> Double

-- | The type for target distributions that can be used in any MCMC sampler.
newtype Target a = T {viewTarget :: TargetView a}
newtype TargetView a = Target (Density a)

-- | Method for constructing custom target distributions.
-- 
-- Target distributions need only a density method.
makeTarget :: Density a -> Target a
makeTarget = T . Target

-- | A procedure that, given a source of randomness, returns an action that 
-- produces a sample. The type itself is read as a verb, i.e, "to sample".
type Sample a = Rand -> IO a

-- | The type for proposal distributions that can be used in any MCMC 
-- sampler.
newtype Proposal a = P {viewProposal :: ProposalView a}
data ProposalView a = Proposal (Density a) (Sample a)

-- | Method for constructing custom proposal distributions.
-- 
-- Proposal distributions need both a density and a sampling method.
makeProposal :: Density a -> Sample a -> Proposal a
makeProposal d = P . Proposal d

-- Kernels --

-- | The type for one step in the random walk. 
-- A value of type @'Step' a@ is a function that takes a source of randomness and a 
-- current state and returns an action producing a subsequent state.
type Step x = Rand -> x -> IO x

-- | The type for MCMC transition kernels. 
-- 
-- The input arguments are the target 
-- distribution (to be modeled) and a /conditional/ proposal distribution.
-- 
-- The result is a 'Step' that will make one move in the random walk based
-- on the current state.
-- In general, an MCMC kernel consists of using:
-- 
-- * the conditioned proposal to make a hypothesis move, and then
-- * the semantics of the specific 'Kernel' at hand to either accept 
-- this hypothesis (and move to the new
-- state), or reject the hypothesis (and stay at the current state).
-- 
-- Type parameter definitions:
-- 
-- [@x@] The kernel-state. This is the type for each state in the /Markov chain/.
-- [@a@] The distribution-state. This is the domain of the target distribution as
-- well as the type of values sampled from the proposal distribution.
-- 
-- In general, we need different types to represent the kernel-state and 
-- distribution-state because the
-- kernel-state may hold extra information that gets updated with each step.
-- Look at 'MCMC.Kernels.simulatedAnnealing' for
-- an example where @x@ differs from @a@.
type Kernel x a = Target a -> (a -> Proposal a) -> Step x

type Act x m a = x -> a -> m a

-- | Type parameter definitions:
-- 
-- [@x@] The kernel-state (see 'MCMC.Types.Kernel')
-- [@a@] The action-state, specific to the action being performed
-- [@m@] The monad in which the action is performed
-- [@b@] The final returned state type
newtype Action x m a b = A {viewAction :: ActionView x m a b}
data ActionView x m a b = Action (Act x m a) (a -> m b) a

makeAction :: Act x m a -- ^ The action to perform at each step of the random walk
           -> (a -> m b) -- ^ The /finish/ function, called at the end of the sampling process
           -> a -- ^ The current value of the action-state
           -> Action x m a b
makeAction act fin = A . (Action act fin)