content: {-# LANGUAGE ImportQualifiedPost #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE TypeFamilies #-}
  
  module ConjugatePriors where
  
  import Control.Applicative (Applicative (..))
  import Control.Foldl (fold)
  import Control.Foldl qualified as F
  import Control.Monad.Bayes.Class (Bayesian (..), MonadDistribution (bernoulli, beta, gamma, normal), MonadMeasure, normalPdf)
  import Numeric.Log (Log (Exp))
  -- Prelude exports liftA2 from GHC 9.6 on, see https://github.com/haskell/core-libraries-committee/blob/main/guides/export-lifta2-prelude.md
  -- import Control.Applicative further up can be removed once we don't support GHC <= 9.4 anymore
  
  import Prelude hiding (Applicative (..))
  
  type GammaParams = (Double, Double)
  
  type BetaParams = (Double, Double)
  
  type NormalParams = (Double, Double)
  
  -- | Posterior on the precision of the normal after the points are observed
  gammaNormalAnalytic ::
    (MonadMeasure m, Foldable t, Functor t) =>
    GammaParams ->
    t Double ->
    m Double
  
  -- | Exact posterior for the model.
  -- For derivation see Kevin Murphy's
  -- "Conjugate Bayesian analysis of the Gaussian distribution"
  -- section 4.
  gammaNormalAnalytic (a, b) points = gamma a' (recip b')
    where
      a' = a + fromIntegral (length points) / 2
      b' = b + sum (fmap (** 2) points) / 2
  
  -- | Posterior on beta after the bernoulli sample
  betaBernoulliAnalytic :: (MonadMeasure m, Foldable t) => BetaParams -> t Bool -> m Double
  betaBernoulliAnalytic (a, b) points = beta a' b'
    where
      (n, s) = fold (liftA2 (,) F.length (F.premap (\case True -> 1; False -> 0) F.sum)) points
      a' = a + s
      b' = b + fromIntegral n - s
  
  bernoulliPdf :: (Floating a) => a -> Bool -> Log a
  bernoulliPdf p x = let numBool = if x then 1.0 else 0 in Exp $ log (p ** numBool * (1 - p) ** (1 - numBool))
  
  betaBernoulli' :: (MonadMeasure m) => (Double, Double) -> Bayesian m Double Bool
  betaBernoulli' (a, b) = Bayesian (beta a b) bernoulli bernoulliPdf
  
  normalNormal' :: (MonadMeasure m) => Double -> (Double, Double) -> Bayesian m Double Double
  normalNormal' var (mu0, var0) = Bayesian (normal mu0 (sqrt var0)) (`normal` (sqrt var)) (`normalPdf` (sqrt var))
  
  gammaNormal' :: (MonadMeasure m) => (Double, Double) -> Bayesian m Double Double
  gammaNormal' (a, b) = Bayesian (gamma a (recip b)) (normal 0 . sqrt . recip) (normalPdf 0 . sqrt . recip)
  
  normalNormalAnalytic ::
    (MonadMeasure m, Foldable t) =>
    Double ->
    NormalParams ->
    t Double ->
    m Double
  normalNormalAnalytic sigma_2 (mu0, sigma0_2) points = normal mu' (sqrt sigma_2')
    where
      (n, s) = fold (liftA2 (,) F.length F.sum) points
      mu' = sigma_2' * (mu0 / sigma0_2 + s / sigma_2)
      sigma_2' = recip (recip sigma0_2 + fromIntegral n / sigma_2)
  