{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}

-- |
-- Module      : Control.Monad.Bayes.Density.Free
-- Description : Free monad transformer over random sampling
-- Copyright   : (c) Adam Scibior, 2015-2020
-- License     : MIT
-- Maintainer  : leonhard.markert@tweag.io
-- Stability   : experimental
-- Portability : GHC
--
-- 'DensityT' is a free monad transformer over random sampling.
module Control.Monad.Bayes.Density.Free
  ( DensityT (..),
    hoist,
    interpret,
    withRandomness,
    runDensityT,
    traced,
  )
where

import Control.Monad.Bayes.Class (MonadDistribution (random))
import Control.Monad.RWS
import Control.Monad.State (evalStateT)
import Control.Monad.Trans.Free.Church (FT, MonadFree (..), hoistFT, iterT, iterTM, liftF)
import Control.Monad.Writer (WriterT (..))
import Data.Functor.Identity (Identity, runIdentity)

-- | Random sampling functor.
newtype SamF a = Random (Double -> a) deriving (Functor)

-- | Free monad transformer over random sampling.
--
-- Uses the Church-encoded version of the free monad for efficiency.
newtype DensityT m a = DensityT {getDensityT :: FT SamF m a}
  deriving newtype (Functor, Applicative, Monad, MonadTrans)

instance MonadFree SamF (DensityT m) where
  wrap = DensityT . wrap . fmap getDensityT

instance (Monad m) => MonadDistribution (DensityT m) where
  random = DensityT $ liftF (Random id)

-- | Hoist 'DensityT' through a monad transform.
hoist :: (Monad m, Monad n) => (forall x. m x -> n x) -> DensityT m a -> DensityT n a
hoist f (DensityT m) = DensityT (hoistFT f m)

-- | Execute random sampling in the transformed monad.
interpret :: (MonadDistribution m) => DensityT m a -> m a
interpret (DensityT m) = iterT f m
  where
    f (Random k) = random >>= k

-- | Execute computation with supplied values for random choices.
withRandomness :: (Monad m) => [Double] -> DensityT m a -> m a
withRandomness randomness (DensityT m) = evalStateT (iterTM f m) randomness
  where
    f (Random k) = do
      xs <- get
      case xs of
        [] -> error "DensityT: the list of randomness was too short"
        y : ys -> put ys >> k y

-- | Execute computation with supplied values for a subset of random choices.
-- Return the output value and a record of all random choices used, whether
-- taken as input or drawn using the transformed monad.
runDensityT :: (MonadDistribution m) => [Double] -> DensityT m a -> m (a, [Double])
runDensityT randomness (DensityT m) =
  runWriterT $ evalStateT (iterTM f $ hoistFT lift m) randomness
  where
    f (Random k) = do
      -- This block runs in StateT [Double] (WriterT [Double]) m.
      -- StateT propagates consumed randomness while WriterT records
      -- randomness used, whether old or new.
      xs <- get
      x <- case xs of
        [] -> random
        y : ys -> put ys >> return y
      tell [x]
      k x

-- | Like 'density', but use an arbitrary sampling monad.
traced :: (MonadDistribution m) => [Double] -> DensityT Identity a -> m (a, [Double])
traced randomness m = runDensityT randomness $ hoist (return . runIdentity) m