content: {-# LANGUAGE RankNTypes #-}
  
  -- |
  -- Module      : Control.Monad.Bayes.Traced.Dynamic
  -- Description : Distributions on execution traces that can be dynamically frozen
  -- Copyright   : (c) Adam Scibior, 2015-2020
  -- License     : MIT
  -- Maintainer  : leonhard.markert@tweag.io
  -- Stability   : experimental
  -- Portability : GHC
  module Control.Monad.Bayes.Traced.Dynamic
    ( TracedT,
      hoist,
      marginal,
      freeze,
      mhStep,
      mh,
    )
  where
  
  import Control.Monad (join)
  import Control.Monad.Bayes.Class
    ( MonadDistribution (random),
      MonadFactor (..),
      MonadMeasure,
    )
  import Control.Monad.Bayes.Density.Free (DensityT)
  import Control.Monad.Bayes.Traced.Common
    ( Trace (..),
      bind,
      mhTransFree,
      scored,
      singleton,
    )
  import Control.Monad.Bayes.Weighted (WeightedT)
  import Control.Monad.Trans (MonadTrans (..))
  import Data.List.NonEmpty as NE (NonEmpty ((:|)), toList)
  
  -- | A tracing monad where only a subset of random choices are traced and this
  -- subset can be adjusted dynamically.
  newtype TracedT m a = TracedT {runTraced :: m (WeightedT (DensityT m) a, Trace a)}
  
  pushM :: (Monad m) => m (WeightedT (DensityT m) a) -> WeightedT (DensityT m) a
  pushM = join . lift . lift
  
  instance (Monad m) => Functor (TracedT m) where
    fmap f (TracedT c) = TracedT $ do
      (m, t) <- c
      let m' = fmap f m
      let t' = fmap f t
      return (m', t')
  
  instance (Monad m) => Applicative (TracedT m) where
    pure x = TracedT $ pure (pure x, pure x)
    (TracedT cf) <*> (TracedT cx) = TracedT $ do
      (mf, tf) <- cf
      (mx, tx) <- cx
      return (mf <*> mx, tf <*> tx)
  
  instance (Monad m) => Monad (TracedT m) where
    (TracedT cx) >>= f = TracedT $ do
      (mx, tx) <- cx
      let m = mx >>= pushM . fmap fst . runTraced . f
      t <- return tx `bind` (fmap snd . runTraced . f)
      return (m, t)
  
  instance MonadTrans TracedT where
    lift m = TracedT $ fmap ((,) (lift $ lift m) . pure) m
  
  instance (MonadDistribution m) => MonadDistribution (TracedT m) where
    random = TracedT $ fmap ((,) random . singleton) random
  
  instance (MonadFactor m) => MonadFactor (TracedT m) where
    score w = TracedT $ fmap (score w,) (score w >> pure (scored w))
  
  instance (MonadMeasure m) => MonadMeasure (TracedT m)
  
  hoist :: (forall x. m x -> m x) -> TracedT m a -> TracedT m a
  hoist f (TracedT c) = TracedT (f c)
  
  -- | Discard the trace and supporting infrastructure.
  marginal :: (Monad m) => TracedT m a -> m a
  marginal (TracedT c) = fmap (output . snd) c
  
  -- | Freeze all traced random choices to their current values and stop tracing
  -- them.
  freeze :: (Monad m) => TracedT m a -> TracedT m a
  freeze (TracedT c) = TracedT $ do
    (_, t) <- c
    let x = output t
    return (return x, pure x)
  
  -- | A single step of the Trace Metropolis-Hastings algorithm.
  mhStep :: (MonadDistribution m) => TracedT m a -> TracedT m a
  mhStep (TracedT c) = TracedT $ do
    (m, t) <- c
    t' <- mhTransFree m t
    return (m, t')
  
  -- | Full run of the Trace Metropolis-Hastings algorithm with a specified
  -- number of steps.
  mh :: (MonadDistribution m) => Int -> TracedT m a -> m [a]
  mh n (TracedT c) = do
    (m, t) <- c
    let f k
          | k <= 0 = return (t :| [])
          | otherwise = do
              (x :| xs) <- f (k - 1)
              y <- mhTransFree m x
              return (y :| x : xs)
    fmap (map output . NE.toList) (f n)
  