content: {-# LANGUAGE DerivingStrategies #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE RankNTypes #-}
  
  -- |
  -- Module      : Control.Monad.Bayes.Sequential
  -- Description : Suspendable probabilistic computation
  -- Copyright   : (c) Adam Scibior, 2015-2020
  -- License     : MIT
  -- Maintainer  : leonhard.markert@tweag.io
  -- Stability   : experimental
  -- Portability : GHC
  --
  -- 'SequentialT' represents a computation that can be suspended.
  module Control.Monad.Bayes.Sequential.Coroutine
    ( SequentialT,
      suspend,
      finish,
      advance,
      finished,
      hoistFirst,
      hoist,
      sequentially,
      sis,
    )
  where
  
  import Control.Monad.Bayes.Class
    ( MonadDistribution (bernoulli, categorical, random),
      MonadFactor (..),
      MonadMeasure,
    )
  import Control.Monad.Coroutine
    ( Coroutine (..),
      bounce,
      mapMonad,
      pogoStick,
    )
  import Control.Monad.Coroutine.SuspensionFunctors
    ( Await (..),
      await,
    )
  import Control.Monad.Trans (MonadIO, MonadTrans (..))
  import Data.Either (isRight)
  
  -- | Represents a computation that can be suspended at certain points.
  -- The intermediate monadic effects can be extracted, which is particularly
  -- useful for implementation of Sequential Monte Carlo related methods.
  -- All the probabilistic effects are lifted from the transformed monad, but
  -- also `suspend` is inserted after each `factor`.
  newtype SequentialT m a = SequentialT {runSequentialT :: Coroutine (Await ()) m a}
    deriving newtype (Functor, Applicative, Monad, MonadTrans, MonadIO)
  
  extract :: Await () a -> a
  extract (Await f) = f ()
  
  instance (MonadDistribution m) => MonadDistribution (SequentialT m) where
    random = lift random
    bernoulli = lift . bernoulli
    categorical = lift . categorical
  
  -- | Execution is 'suspend'ed after each 'score'.
  instance (MonadFactor m) => MonadFactor (SequentialT m) where
    score w = lift (score w) >> suspend
  
  instance (MonadMeasure m) => MonadMeasure (SequentialT m)
  
  -- | A point where the computation is paused.
  suspend :: (Monad m) => SequentialT m ()
  suspend = SequentialT await
  
  -- | Remove the remaining suspension points.
  finish :: (Monad m) => SequentialT m a -> m a
  finish = pogoStick extract . runSequentialT
  
  -- | Execute to the next suspension point.
  -- If the computation is finished, do nothing.
  --
  -- > finish = finish . advance
  advance :: (Monad m) => SequentialT m a -> SequentialT m a
  advance = SequentialT . bounce extract . runSequentialT
  
  -- | Return True if no more suspension points remain.
  finished :: (Monad m) => SequentialT m a -> m Bool
  finished = fmap isRight . resume . runSequentialT
  
  -- | Transform the inner monad.
  -- This operation only applies to computation up to the first suspension.
  hoistFirst :: (forall x. m x -> m x) -> SequentialT m a -> SequentialT m a
  hoistFirst f = SequentialT . Coroutine . f . resume . runSequentialT
  
  -- | Transform the inner monad.
  -- The transformation is applied recursively through all the suspension points.
  hoist ::
    (Monad m, Monad n) =>
    (forall x. m x -> n x) ->
    SequentialT m a ->
    SequentialT n a
  hoist f = SequentialT . mapMonad f . runSequentialT
  
  -- | Apply a function a given number of times.
  composeCopies :: Int -> (a -> a) -> (a -> a)
  composeCopies k f = foldr (.) id (replicate k f)
  
  -- | Sequential importance sampling.
  -- Applies a given transformation after each time step.
  sequentially,
    sis ::
      (Monad m) =>
      -- | transformation
      (forall x. m x -> m x) ->
      -- | number of time steps
      Int ->
      SequentialT m a ->
      m a
  sequentially f k = finish . composeCopies k (advance . hoistFirst f)
  
  -- | synonym
  sis = sequentially
  