content: {-# LANGUAGE TupleSections #-}
  
  module Bayesian where
  
  import Data.Ratio
  import Data.Tuple (swap)
  
  ------------------------------------------------------------------------
  -- BAYESIAN EXPERIMENTS
  ------------------------------------------------------------------------
  
  -- A Bayesian experiment is a list of values each paired with a rational
  -- describing their respective probability. We expect the sum of these
  -- rationals to be less than 1.
  -- A user should not be able to define such a thing by hand, and she
  -- should not be able to `unsafeRun` it.
  newtype Bayesian a = Bayesian { unsafeRunBayesian :: [(Rational, a)] }
  
  -- If she wants to observe the result of the experiment, she has to use
  -- one of the following functions.
  
  -- `runBayesian` groups together the outcomes which are equal (by summing
  -- their probabilities) and normalises the distribution so that the sum
  -- of rationals is equal to one.
  -- * The grouping is necessary because the same outcome might be
  --   generated by different branches of an experiment.
  -- * The normalisation is necessary because making an observation might
  --   invalidate some runs. This means that the result of an experiment
  --   is a subdistribution composed of the outcomes of the valid runs.
  
  runBayesian :: Eq a => Bayesian a -> [(Rational, a)]
  runBayesian b = reverse xs where
  
    (s, xs) = norm $ unsafeRunBayesian b
    norm    = foldr (\ (p, a) (s', ih) -> (p + s', insert (p / s) a ih)) (0, [])
  
    insert :: Eq a => Rational -> a -> [(Rational, a)] -> [(Rational, a)]
    insert p a [] = [(p, a)]
    insert p a (x@(q, b) : xs)
      | a == b    = (p + q , a) : xs
      | otherwise = x : insert p a xs
  
  -- `checkBayesian` runs the experiment and only retains the value we
  -- are interested in. A typical use-case is `epidemiology` where we
  -- are interested in the probability that the person is indeed sick
  -- (i.e. that the returned `hasDisease` is equal to `True`).
  
  checkBayesian :: Eq a => a -> Bayesian a -> Rational
  checkBayesian a = maybe 0 id . lookup a . fmap swap . runBayesian
  
  -- We can define the expected `Functor`, `Applicative` and `Monad`
  -- instances for `Bayesian`
  instance Functor Bayesian where
    fmap f = Bayesian . fmap (f <$>) . unsafeRunBayesian
  
  instance Applicative Bayesian where
    pure      = Bayesian . pure . (1 % 1,)
    fs <*> ts = Bayesian [ (p * q, f t) | (p, f) <- unsafeRunBayesian fs
                                        , (q, t) <- unsafeRunBayesian ts ]
  
  instance Monad Bayesian where
    return   = pure
    ts >>= f = Bayesian $ concatMap continuation $ unsafeRunBayesian ts where
      continuation (p, t) = adjustBy p <$> next where
        next = unsafeRunBayesian $ f t
        adjustBy p (q, ft) = (p * q, ft)
  
  ------------------------------------------------------------------------
  -- COMBINATORS
  ------------------------------------------------------------------------
  
  -- The classic `bernouilli` experiment: we get `True` with probability
  -- `p` and `False` otherwise.
  bernouilli :: Rational -> Bayesian Bool
  bernouilli p = Bayesian [ (p, True), (1-p, False) ]
  
  -- The usual `uniform` distribution: each value has the same probability
  -- of being selected.
  uniform :: [a] -> Bayesian a
  uniform xs = Bayesian $ (1 % n,) <$> xs
    where n = fromIntegral $ length xs
  
  -- `observe` states that an observation must be `True`. If it is not
  -- then the corresponding run is considered invalid an simply thrown
  -- away.
  observe :: Bool -> Bayesian Bool
  observe b = if b then pure True else Bayesian []
  
  ------------------------------------------------------------------------
  -- EXAMPLES
  ------------------------------------------------------------------------
  
  -- These two examples are taken from
  -- Borgström, Gordon, Greenberg, Margetson and van Gael's
  -- Measure Transformer Semantics for Bayesian Machine Learning
  epidemiology :: Bayesian Bool
  epidemiology = do
    hasDisease     <- bernouilli (1 % 100)
    positiveResult <- if hasDisease
                      then bernouilli (80 % 100)
                      else bernouilli (96 % 1000)
    observe positiveResult
    return hasDisease
  
  noTwoTails :: Bayesian (Bool, Bool)
  noTwoTails = do
    heads1 <- bernouilli (1 % 2)
    heads2 <- bernouilli (1 % 2)
    observe (heads1 || heads2)
    return (heads1, heads2)
  
  -- This example is taken from Bart Jacoband Robin Adams'
  -- A Type Theory for Probabilistic and Bayesian Reasoning.
  alarmRinging :: Bayesian Bool
  alarmRinging = do
    burglar    <- bernouilli (1 % 1000)
    earthquake <- bernouilli (1 % 500)
    alarm      <- rings burglar earthquake
    observe =<< johnCalls alarm
    return burglar
  
    where
  
      johnCalls alarm = bernouilli $ if alarm then 9 % 10 else 1 % 20
      maryCalls alarm = bernouilli $ if alarm then 7 % 10 else 1 % 100
  
      rings b e = bernouilli $ case (b , e) of
        (True , True)  -> 95 % 100
        (True , False) -> 94 % 100
        (False, True)  -> 29 % 100
        (False, False) -> 1 % 1000
  