DiscreetDistribution.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
module Data.DiscreetDistribution where

import Control.Monad.Trans.Free
import Control.Monad


newtype DiscreetDistribution a = DiscreetDistribution {
    unDD :: [(a, Double)]
    }


instance Show a => Show (DiscreetDistribution a) where
    show (DiscreetDistribution dist) = "[\n" ++ (unlines $ map show dist) ++ "]\n"


data SamplerF a b
    = Sample (a -> b)


instance Functor (SamplerF a) where
    fmap f (Sample c) = Sample $ f . c


type Sampler a = Free (SamplerF a)
type SamplferT a = FreeT (SamplerF a)


sample :: MonadFree (SamplerF a) m => m a
sample = liftF $ Sample id


runSamplerT
    :: Monad m
    => SamplerT a m ()
    -> DiscreetDistribution a
    -> m (DiscreetDistribution [a])
runSamplerT d dist = runFreeT d >>= \case
    Pure _ -> return $ DiscreetDistribution [([], 1)]
    Free (Sample cont) -> undefined -- What do I do here???


runSampler
    :: Sampler a ()
    -> DiscreetDistribution a
    -> DiscreetDistribution [a]
runSampler d dist = case runFree d of
    Pure _ -> DiscreetDistribution $ [([], 1)]
    Free (Sample cont) -> DiscreetDistribution $ do
        (val, prob) <-  unDD dist
        let futurePulls = cont val
        (futurePulledVals, futurePulledProb) <- unDD $ runSampler futurePulls dist

        return $ (val:futurePulledVals, prob * futurePulledProb)


data Ball = Red | Green | Blue deriving Show


testDist :: DiscreetDistribution Ball
testDist = DiscreetDistribution $ [(Red, 0.3), (Green, 0.5), (Blue, 0.2)]


testDistExperiment :: Sampler Ball ()
testDistExperiment = do
    first <- sample
    case first of
        Red -> return ()
        _ -> do
            void sample
            void sample
            return ()


resultDist :: DiscreetDistribution [Ball]
resultDist = runSampler testDistExperiment testDist


main :: IO ()
main = print resultDist