{-# LANGUAGE ImportQualifiedPost #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}

module BetaBin where

-- The beta-binomial model in latent variable and urn model representations.
-- The two formulations should be exactly equivalent, but only urn works with Dist.
import Control.Monad (replicateM)
import Control.Monad.Bayes.Class
  ( MonadDistribution (bernoulli, uniform),
    MonadMeasure,
    condition,
  )
import Control.Monad.State.Lazy (evalStateT, get, put)
import Pipes ((<-<))
import Pipes.Prelude qualified as P hiding (show)

-- | Beta-binomial model as an i.i.d. sequence conditionally on weight.
latent :: MonadDistribution m => Int -> m [Bool]
latent n = do
  weight <- uniform 0 1
  replicateM n (bernoulli weight)

-- | Beta-binomial as a random process.
-- Equivalent to the above by De Finetti's theorem.
urn :: MonadDistribution m => Int -> m [Bool]
urn n = flip evalStateT (1, 1) $ do
  replicateM n do
    (a, b) <- get
    let weight = a / (a + b)
    outcome <- bernoulli weight
    let (a', b') = if outcome then (a + 1, b) else (a, b + 1)
    put (a', b')
    return outcome

-- | Beta-binomial as a random process.
-- This time using the Pipes library, for a more pure functional style
urnP :: MonadDistribution m => Int -> m [Bool]
urnP n = P.toListM $ P.take n <-< P.unfoldr toss (1, 1)
  where
    toss (a, b) = do
      let weight = a / (a + b)
      outcome <- bernoulli weight
      let (a', b') = if outcome then (a + 1, b) else (a, b + 1)
      return $ Right (outcome, (a', b'))

-- | A beta-binomial model where the first three states are True,True,False.
-- The resulting distribution is on the remaining outcomes.
cond :: MonadMeasure m => m [Bool] -> m [Bool]
cond d = do
  ~(first : second : third : rest) <- d
  condition first
  condition second
  condition (not third)
  return rest

-- | The final conditional model, abstracting the representation.
model :: MonadMeasure m => (Int -> m [Bool]) -> Int -> m Int
model repr n = fmap count $ cond $ repr (n + 3)
  where
    -- Post-processing by counting the number of True values.
    count :: [Bool] -> Int
    count = length . filter id