content: {-# LANGUAGE ImportQualifiedPost #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# OPTIONS_GHC -Wno-name-shadowing #-}
  
  module Control.Monad.Bayes.Inference.Lazy.MH where
  
  import Control.Monad.Bayes.Class (Log (ln))
  import Control.Monad.Bayes.Sampler.Lazy
    ( Sampler,
      Tree (..),
      Trees (..),
      randomTree,
      runSampler,
    )
  import Control.Monad.Bayes.Weighted (WeightedT, runWeightedT)
  import Control.Monad.Extra (iterateM)
  import Control.Monad.State.Lazy (MonadState (get, put), runState)
  import System.Random (RandomGen (split), getStdGen, newStdGen)
  import System.Random qualified as R
  
  mh :: forall a. Double -> WeightedT Sampler a -> IO [(a, Log Double)]
  mh p m = do
    -- Top level: produce a stream of samples.
    -- Split the random number generator in two
    -- One part is used as the first seed for the simulation,
    -- and one part is used for the randomness in the MH algorithm.
    g <- newStdGen >> getStdGen
    let (g1, g2) = split g
    let t = randomTree g1
    let (x, w) = runSampler (runWeightedT m) t
    -- Now run step over and over to get a stream of (tree,result,weight)s.
    let (samples, _) = runState (iterateM step (t, x, w)) g2
    -- The stream of seeds is used to produce a stream of result/weight pairs.
    return $ map (\(_, x, w) -> (x, w)) samples
    where
      --   where
      {- NB There are three kinds of randomness in the step function.
      1. The start tree 't', which is the source of randomness for simulating the
      program m to start with. This is sort-of the point in the "state space".
      2. The randomness needed to propose a new tree ('g1')
      3. The randomness needed to decide whether to accept or reject that ('g2')
      The tree t is an argument and result,
      but we use a state monad ('get'/'put') to deal with the other randomness '(g,g1,g2)' -}
  
      -- step :: RandomGen g => (Tree, a, Log Double) -> State g (Tree, a, Log Double)
      step (t, x, w) = do
        -- Randomly change some sites
        g <- get
        let (g1, g2) = split g
        let t' = mutateTree p g1 t
        -- Rerun the model with the new tree, to get a new
        -- weight w'.
        let (x', w') = runSampler (runWeightedT m) t'
        -- MH acceptance ratio. This is the probability of either
        -- returning the new seed or the old one.
        let ratio = w' / w
        let (r, g2') = R.random g2
        put g2'
        if r < min 1 (exp $ ln ratio)
          then return (t', x', w')
          else return (t, x, w)
  
  -- Replace the labels of a tree randomly, with probability p
  mutateTree :: forall g. (RandomGen g) => Double -> g -> Tree -> Tree
  mutateTree p g (Tree a ts) =
    let (a', g') = (R.random g :: (Double, g))
        (a'', g'') = R.random g'
     in Tree
          { currentUniform = if a' < p then a'' else a,
            lazyUniforms = mutateTrees p g'' ts
          }
  
  mutateTrees :: (RandomGen g) => Double -> g -> Trees -> Trees
  mutateTrees p g (Trees t ts) =
    let (g1, g2) = split g
     in Trees
          { headTree = mutateTree p g1 t,
            tailTrees = mutateTrees p g2 ts
          }
  