content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE DeriveFunctor #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  -- | This is a port of the implementation of LazyPPL: https://lazyppl.bitbucket.io/
  module Control.Monad.Bayes.Sampler.Lazy where
  
  import Control.Monad (ap)
  import Control.Monad.Bayes.Class (MonadDistribution (random))
  import Control.Monad.Bayes.Weighted (WeightedT, runWeightedT)
  import Control.Monad.IO.Class
  import Control.Monad.Identity (Identity (runIdentity))
  import Control.Monad.Trans
  import Numeric.Log (Log (..))
  import System.Random
    ( RandomGen (split),
      getStdGen,
      newStdGen,
    )
  import System.Random qualified as R
  
  -- | A 'Tree' is a lazy, infinitely wide and infinitely deep tree, labelled by Doubles.
  --
  --   Our source of randomness will be a Tree, populated by uniform [0,1] choices for each label.
  --   Often people just use a list or stream instead of a tree.
  --   But a tree allows us to be lazy about how far we are going all the time.
  data Tree = Tree
    { currentUniform :: Double,
      lazyUniforms :: Trees
    }
  
  -- | An infinite stream of 'Tree's.
  data Trees = Trees
    { headTree :: Tree,
      tailTrees :: Trees
    }
  
  -- | A probability distribution over @a@ is a function 'Tree -> a'.
  --   The idea is that it uses up bits of the tree as it runs.
  type Sampler = SamplerT Identity
  
  runSampler :: Sampler a -> Tree -> a
  runSampler = (runIdentity .) . runSamplerT
  
  newtype SamplerT m a = SamplerT {runSamplerT :: Tree -> m a}
    deriving (Functor)
  
  -- | Split a tree in two (bijectively).
  splitTree :: Tree -> (Tree, Tree)
  splitTree (Tree r (Trees t ts)) = (t, Tree r ts)
  
  -- | Generate a tree with uniform random labels.
  --
  -- Preliminary for the simulation methods. This uses 'split' to split a random seed.
  randomTree :: (RandomGen g) => g -> Tree
  randomTree g = let (a, g') = R.random g in Tree a (randomTrees g')
  
  randomTrees :: (RandomGen g) => g -> Trees
  randomTrees g = let (g1, g2) = split g in Trees (randomTree g1) (randomTrees g2)
  
  instance (Monad m) => Applicative (SamplerT m) where
    pure = lift . pure
    (<*>) = ap
  
  -- | Sequencing is done by splitting the tree
  --   and using different bits for different computations.
  instance (Monad m) => Monad (SamplerT m) where
    return = pure
    (SamplerT m) >>= f = SamplerT \g -> do
      let (g1, g2) = splitTree g
      a <- m g1
      let SamplerT m' = f a
      m' g2
  
  instance MonadTrans SamplerT where
    lift = SamplerT . const
  
  instance (MonadIO m) => MonadIO (SamplerT m) where
    liftIO = lift . liftIO
  
  -- | Sampling gets the label at the head of the tree and discards the rest.
  instance (Monad m) => MonadDistribution (SamplerT m) where
    random = SamplerT \(Tree r _) -> pure r
  
  -- | Runs a 'SamplerT' by creating a new 'StdGen'.
  runSamplerTIO :: (MonadIO m) => SamplerT m a -> m a
  runSamplerTIO m = liftIO newStdGen *> (runSamplerT m =<< randomTree <$> liftIO getStdGen)
  
  -- | Draw a stream of independent samples.
  independent :: (Monad m) => m a -> m [a]
  independent = sequence . repeat
  
  -- | Runs a probability measure and gets out a stream of @(result,weight)@ pairs
  weightedSamples :: (MonadIO m) => WeightedT (SamplerT m) a -> m [(a, Log Double)]
  weightedSamples = runSamplerTIO . sequence . repeat . runWeightedT
  