content: {-# LANGUAGE DerivingStrategies #-}
  {-# LANGUAGE ImportQualifiedPost #-}
  
  module Helper where
  
  import Control.Monad.Bayes.Class (MonadMeasure)
  import Control.Monad.Bayes.Inference.MCMC (MCMCConfig (..), Proposal (SingleSiteMH))
  import Control.Monad.Bayes.Inference.RMSMC (rmsmcBasic)
  import Control.Monad.Bayes.Inference.SMC
    ( SMCConfig (SMCConfig, numParticles, numSteps, resampler),
      smc,
    )
  import Control.Monad.Bayes.Population
  import Control.Monad.Bayes.Sampler.Strict
  import Control.Monad.Bayes.Traced hiding (model)
  import Control.Monad.Bayes.Weighted
  import Control.Monad.ST (runST)
  import HMM qualified
  import LDA qualified
  import LogReg qualified
  
  data Model = LR Int | HMM Int | LDA (Int, Int)
    deriving stock (Show, Read)
  
  parseModel :: String -> Maybe Model
  parseModel s =
    case s of
      'L' : 'R' : n -> Just $ LR (read n)
      'H' : 'M' : 'M' : n -> Just $ HMM (read n)
      'L' : 'D' : 'A' : n -> Just $ LDA (5, read n)
      _ -> Nothing
  
  serializeModel :: Model -> Maybe String
  serializeModel (LR n) = Just $ "LR" ++ show n
  serializeModel (HMM n) = Just $ "HMM" ++ show n
  serializeModel (LDA (5, n)) = Just $ "LDA" ++ show n
  serializeModel (LDA _) = Nothing
  
  data Alg = SMC | MH | RMSMC
    deriving stock (Read, Show, Eq, Ord, Enum, Bounded)
  
  getModel :: (MonadMeasure m) => Model -> (Int, m String)
  getModel model = (size model, program model)
    where
      size (LR n) = n
      size (HMM n) = n
      size (LDA (d, w)) = d * w
      program (LR n) = show <$> (LogReg.logisticRegression (runST $ sampleSTfixed (LogReg.syntheticData n)))
      program (HMM n) = show <$> (HMM.hmm (runST $ sampleSTfixed (HMM.syntheticData n)))
      program (LDA (d, w)) = show <$> (LDA.lda (runST $ sampleSTfixed (LDA.syntheticData d w)))
  
  runAlg :: Model -> Alg -> SamplerIO String
  runAlg model alg =
    case alg of
      SMC ->
        let n = 100
            (k, m) = getModel model
         in show <$> runPopulationT (smc SMCConfig {numSteps = k, numParticles = n, resampler = resampleSystematic} m)
      MH ->
        let t = 100
            (_, m) = getModel model
         in show <$> unweighted (mh t m)
      RMSMC ->
        let n = 10
            t = 1
            (k, m) = getModel model
         in show <$> runPopulationT (rmsmcBasic MCMCConfig {numMCMCSteps = t, numBurnIn = 0, proposal = SingleSiteMH} (SMCConfig {numSteps = k, numParticles = n, resampler = resampleSystematic}) m)
  
  runAlgFixed :: Model -> Alg -> IO String
  runAlgFixed model alg = sampleIOfixed $ runAlg model alg
  