module CS4400.Choice where

import Control.Monad (ap, liftM)
import System.Random (StdGen, uniformR)
import System.Random.Stateful (initStdGen)

class MonadChoice m where
  choose :: [a] -> m a

instance MonadChoice [] where
  choose = id

data Sample a = Sample {runSample :: StdGen -> (StdGen, a)}

instance Monad Sample where
  m >>= f =
    Sample
      ( \g ->
          let (g', x) = runSample m g
           in runSample (f x) g'
      )

instance MonadChoice Sample where
  choose xs =
    Sample
      ( \g ->
          let (i, g') = uniformR (0, length xs - 1) g
           in (g', xs !! i)
      )

instance Applicative Sample where
  pure x = Sample (\g -> (g, x))
  (<*>) = ap

instance Functor Sample where
  fmap = liftM

sample :: Sample String -> IO String
sample s = do
  g <- initStdGen
  pure (snd (runSample s g))

genstr :: (Monad m, MonadChoice m) => Int -> String -> String -> m String
genstr len alphabet suffix =
  if len <= 0
    then pure suffix
    else do
      c <- choose alphabet
      genstr (len - 1) alphabet (c : suffix)