-- | A general purpose library for simulating differential processes, of a deterministic or
-- stochastic nature.

module Goal.Simulation.Chain (
    -- * Chains
      Chain
    , streamChain
    -- ** Markov Chains
    , MarkovTensor
    , markovTensor
    , markovChain
    -- ** Harris Chains
    , harrisChain
    ) where

--- Imports ---


-- Goal --

import Goal.Simulation.Mealy

import Goal.Geometry
import Goal.Probability


--- Stochastic Processes ---


-- | A 'Chain' is a discrete time 'Flow'.
type Chain s = Mealy () s

streamChain :: Chain k -> [k]
-- | A convenience function for streaming 'MarkovChain's.
streamChain mchn = stream mchn $ repeat ()

-- Finite --

type MarkovTensor s = Tensor (CurvedCategorical s) (CurvedCategorical s)

markovTensor :: s -> MarkovTensor s
markovTensor s = Tensor (CurvedCategorical s) (CurvedCategorical s)

markovChain :: Discrete s
    => (Function Standard Standard :#: MarkovTensor s) -- ^ The stochastic matrix
    -> Element s -- ^ The initial state
    -> RandST r (Chain (Element s)) -- ^ An embedded markov chain
-- | Constructs a markov chain process.
markovChain smtx =
    accumulateRandomFunction (markovChainAccumulator smtx)

markovChainAccumulator smtx _ k = do
    let ks = elements . sampleSpace . domain $ manifold smtx
        sk = fromList (domain $ manifold smtx) [ if k' == k then 1 else 0 | k' <- ks ]
    k' <- generate $ smtx >.> sk
    return (k',k')

-- Continuous --

harrisChain :: (Manifold m, Transition c Standard m, Generative Standard m)
    => (Sample m -> c :#: m)
    -> Sample m
    -> RandST s' (Chain (Sample m))
harrisChain mdl =
    accumulateRandomFunction (harrisChainAccumulator mdl)

harrisChainAccumulator mdl () s = do
    let cp' = mdl s
    s' <- standardGenerate cp'
    return (s',s')