content: -- HMM from Anglican (https://bitbucket.org/probprog/anglican-white-paper)
  
  module HMM where
  
  import Control.Monad (replicateM, when)
  import Control.Monad.Bayes.Class
    ( MonadDistribution (categorical, normal, uniformD),
      MonadFactor,
      MonadMeasure,
      factor,
      normalPdf,
    )
  import Control.Monad.Bayes.Enumerator (enumerateToDistribution)
  import Data.Maybe (fromJust, isJust)
  import Data.Vector (fromList)
  import Pipes (MFunctor (hoist), MonadTrans (lift), each, yield, (>->))
  import Pipes.Core (Producer)
  import Pipes.Prelude qualified as Pipes
  
  -- | Observed values
  values :: [Double]
  values =
    [ 0.9,
      0.8,
      0.7,
      0,
      -0.025,
      -5,
      -2,
      -0.1,
      0,
      0.13,
      0.45,
      6,
      0.2,
      0.3,
      -1,
      -1
    ]
  
  -- | The transition model.
  trans :: (MonadDistribution m) => Int -> m Int
  trans 0 = categorical $ fromList [0.1, 0.4, 0.5]
  trans 1 = categorical $ fromList [0.2, 0.6, 0.2]
  trans 2 = categorical $ fromList [0.15, 0.7, 0.15]
  trans _ = error "unreachable"
  
  -- | The emission model.
  emissionMean :: Int -> Double
  emissionMean 0 = -1
  emissionMean 1 = 1
  emissionMean 2 = 0
  emissionMean _ = error "unreachable"
  
  -- | Initial state distribution
  start :: (MonadDistribution m) => m Int
  start = uniformD [0, 1, 2]
  
  -- | Example HMM from http://dl.acm.org/citation.cfm?id=2804317
  hmm :: (MonadMeasure m) => [Double] -> m [Int]
  hmm dataset = f dataset (const . return)
    where
      expand x y = do
        x' <- trans x
        factor $ normalPdf (emissionMean x') 1 y
        return x'
      f [] k = start >>= k []
      f (y : ys) k = f ys (\xs x -> expand x y >>= k (x : xs))
  
  syntheticData :: (MonadDistribution m) => Int -> m [Double]
  syntheticData n = replicateM n syntheticPoint
    where
      syntheticPoint = uniformD [0, 1, 2]
  
  -- | Equivalent model, but using pipes for simplicity
  
  -- | Prior expressed as a stream
  hmmPrior :: (MonadDistribution m) => Producer Int m b
  hmmPrior = do
    x <- lift start
    yield x
    Pipes.unfoldr (fmap (Right . (\k -> (k, k))) . trans) x
  
  -- | Observations expressed as a stream
  hmmObservations :: (Functor m) => [a] -> Producer (Maybe a) m ()
  hmmObservations dataset = each (Nothing : (Just <$> reverse dataset))
  
  -- | Posterior expressed as a stream
  hmmPosterior :: (MonadMeasure m) => [Double] -> Producer Int m ()
  hmmPosterior dataset =
    zipWithM
      hmmLikelihood
      hmmPrior
      (hmmObservations dataset)
    where
      hmmLikelihood :: (MonadFactor f) => (Int, Maybe Double) -> f ()
      hmmLikelihood (l, o) = when (isJust o) (factor $ normalPdf (emissionMean l) 1 (fromJust o))
  
      zipWithM f p1 p2 = Pipes.zip p1 p2 >-> Pipes.chain f >-> Pipes.map fst
  
  hmmPosteriorPredictive :: (MonadDistribution m) => [Double] -> Producer Double m ()
  hmmPosteriorPredictive dataset =
    Pipes.hoist enumerateToDistribution (hmmPosterior dataset)
      >-> Pipes.mapM (\x -> normal (emissionMean x) 1)
  