{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE BangPatterns #-}

module PPL.Sampling where

import Control.Monad.IO.Class
import Control.Monad.Trans.State
import Data.Bifunctor
import Data.Monoid
import Numeric.Log
import PPL.Distr
import PPL.Internal
import System.Random (StdGen, random, randoms)
import qualified Streaming as S
import Streaming.Prelude (Stream, yield, Of)

mh :: Monad m => StdGen -> Double -> Meas a -> Stream (Of (a, Log Double)) m ()
mh g p m = step t0 t x w
  where
    (t0,t) = split $ randomTree g
    (x, w) = head $ samples m t

    step !t0 !t !x !w = do
      let (t1:t2:t3:t4:_) = splitTrees t0
          t' = mutateTree p t1 t2 t
          (x', w') = head $ samples m t'
          ratio = w' / w
          (Exp . log -> r) = draw t3
          (t'', x'', w'') = if r < ratio
            then (t', x', w')
            else (t, x, w)
      yield (x'', w'')
      step t4 t'' x'' w''