module TestSequential (twoSync, finishedTwoSync, checkTwoSync, checkPreserve, pFinished, isFinished, checkSync) where

import Control.Monad.Bayes.Class
  ( MonadDistribution (uniformD),
    MonadMeasure,
    factor,
  )
import Control.Monad.Bayes.Enumerator as Dist (enumerator, mass)
import Control.Monad.Bayes.Sequential.Coroutine (advance, finish, finished)
import Data.AEq (AEq ((~==)))
import Sprinkler (soft)

twoSync :: (MonadMeasure m) => m Int
twoSync = do
  x <- uniformD [0, 1]
  factor (fromIntegral x)
  y <- uniformD [0, 1]
  factor (fromIntegral y)
  return (x + y)

finishedTwoSync :: (MonadMeasure m) => Int -> m Bool
finishedTwoSync n = finished (run n twoSync)
  where
    run 0 d = d
    run k d = run (k - 1) (advance d)

checkTwoSync :: Int -> Bool
checkTwoSync 0 = mass (finishedTwoSync 0) False ~== 1
checkTwoSync 1 = mass (finishedTwoSync 1) False ~== 1
checkTwoSync 2 = mass (finishedTwoSync 2) True ~== 1
checkTwoSync _ = error "Unexpected argument"

sprinkler :: (MonadMeasure m) => m Bool
sprinkler = Sprinkler.soft

checkPreserve :: Bool
checkPreserve = enumerator (finish sprinkler) ~== enumerator sprinkler

pFinished :: Int -> Double
pFinished 0 = 0.8267716535433071
pFinished 1 = 0.9988062077198566
pFinished 2 = 1
pFinished _ = error "Unexpected argument"

isFinished :: (MonadMeasure m) => Int -> m Bool
isFinished n = finished (run n sprinkler)
  where
    run 0 d = d
    run k d = run (k - 1) (advance d)

checkSync :: Int -> Bool
checkSync n = mass (isFinished n) True ~== pFinished n