-- | Effectful samplers as @Prob (K IO)@ morphisms.
module Circuit.Inference.Sampler
  ( -- * Primitive samplers
    bernoulli,
    uniformDiscrete,

    -- * Recursive trace samplers
    geometric,
    geometricBody,
    sample,
  )
where

import Circuit.Category (K (..))
import Circuit.Inference.Prob (Prob (..), fromWeightedK, traceEK)
import System.Random (randomRIO)

-- | Bernoulli trial with success probability @p@.
--
-- Output 'True' with probability @p@, 'False' with probability @1-p@.
bernoulli :: Double -> Prob (K IO) Double () Bool
bernoulli :: Double -> Prob (K IO) Double () Bool
bernoulli Double
p = [(Bool, Double)] -> Prob (K IO) Double () Bool
forall r (m :: * -> *) b.
(Num r, Monad m) =>
[(b, r)] -> Prob (K m) r () b
fromWeightedK [(Bool
True, Double
p), (Bool
False, Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
p)]

-- | Discrete uniform distribution over a finite list.
uniformDiscrete :: [a] -> Prob (K IO) Double () a
uniformDiscrete :: forall a. [a] -> Prob (K IO) Double () a
uniformDiscrete [a]
xs =
  let n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
      w :: Double
w = Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
   in [(a, Double)] -> Prob (K IO) Double () a
forall r (m :: * -> *) b.
(Num r, Monad m) =>
[(b, r)] -> Prob (K m) r () b
fromWeightedK [(a
x, Double
w) | a
x <- [a]
xs]

-- | Extract a single sample from a sampler by using the output itself as the
-- dualizing object.
sample :: Prob (K IO) b () b -> IO b
sample :: forall b. Prob (K IO) b () b -> IO b
sample Prob (K IO) b () b
s = K IO ((), ()) b -> ((), ()) -> IO b
forall {k} (m :: k -> *) a (b :: k). K m a b -> a -> m b
runK (Prob (K IO) b () b -> forall x. K IO (x, b) b -> K IO (x, ()) b
forall {k} (arr :: * -> k -> *) (r :: k) a b.
Prob arr r a b -> forall x. arr (x, b) r -> arr (x, a) r
runProb Prob (K IO) b () b
s ((((), b) -> IO b) -> K IO ((), b) b
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K (\(()
_, b
b) -> b -> IO b
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
b))) ((), ())

-- | Geometric distribution counting failures before the first success.
--
-- Implemented as a terminating 'traceEK' over an effectful Bernoulli body.
-- Each recursive iteration performs an 'IO' sample, so the trace is productive
-- and returns a value with probability 1 (for @0 < p <= 1@).
geometric :: Double -> Prob (K IO) Int () Int
geometric :: Double -> Prob (K IO) Int () Int
geometric Double
p = Prob (K IO) Int (Either () Int) (Either Int Int)
-> Prob (K IO) Int () Int
forall {k} (m :: k -> *) (r :: k) a s b.
Prob (K m) r (Either a s) (Either b s) -> Prob (K m) r a b
traceEK (Double -> Prob (K IO) Int (Either () Int) (Either Int Int)
geometricBody Double
p)

-- | Body of the geometric sampler.
--
-- Input: @Left ()@ on first call, @Right n@ on recursive calls where @n@ is
-- the current failure count.  Output: @Right k@ terminates with final count
-- @k@; @Left k@ continues with count @k@.
geometricBody ::
  Double ->
  Prob (K IO) Int (Either () Int) (Either Int Int)
geometricBody :: Double -> Prob (K IO) Int (Either () Int) (Either Int Int)
geometricBody Double
p = (forall x.
 K IO (x, Either Int Int) Int -> K IO (x, Either () Int) Int)
-> Prob (K IO) Int (Either () Int) (Either Int Int)
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x.
  K IO (x, Either Int Int) Int -> K IO (x, Either () Int) Int)
 -> Prob (K IO) Int (Either () Int) (Either Int Int))
-> (forall x.
    K IO (x, Either Int Int) Int -> K IO (x, Either () Int) Int)
-> Prob (K IO) Int (Either () Int) (Either Int Int)
forall a b. (a -> b) -> a -> b
$ \(K (x, Either Int Int) -> IO Int
k) -> ((x, Either () Int) -> IO Int) -> K IO (x, Either () Int) Int
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K (((x, Either () Int) -> IO Int) -> K IO (x, Either () Int) Int)
-> ((x, Either () Int) -> IO Int) -> K IO (x, Either () Int) Int
forall a b. (a -> b) -> a -> b
$ \(x
x, Either () Int
e) -> do
  u <- (Double, Double) -> IO Double
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Double
0, Double
1) :: IO Double
  let out = case Either () Int
e of
        -- Left value escapes as the final output; Right value feeds back.
        -- State counts failures so far, so the first failure advances to 1.
        Left () -> if Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
p then Int -> Either Int Int
forall a b. a -> Either a b
Left Int
0 else Int -> Either Int Int
forall a b. b -> Either a b
Right Int
1
        Right Int
n -> if Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
p then Int -> Either Int Int
forall a b. a -> Either a b
Left Int
n else Int -> Either Int Int
forall a b. b -> Either a b
Right (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  k (x, out)