Fixed bugs in simple tri-state Markov Chain generator

Markov.hs

#!/usr/bin/env stack
-- stack --resolver lts-13.30 script --package random

import           Control.Monad (foldM)
import           Data.List     (elemIndex, minimumBy)
import           Data.Maybe    (fromMaybe)
import           Debug.Trace   (trace)
import           System.IO     (getLine)
import           System.Random (randomIO, randomRIO)

data State = A | B | C
  deriving (Eq, Read, Show)

type StateF = (State, [Double])

type Table = [StateF]

-- | The mapping from a state to the probability of the other states
table :: Table
table = [(A, [0.3, 0.5, 0.2])
        ,(B, [0.7, 0.3, 0.0])
        ,(C, [0.2, 0.1, 0.7])]

-- | Find the next state in @Table@ according to the current @State@ and the given probability
step :: State -> Table -> Double -> State
step cs t x = do
  let r = fst $ t !! ns
  trace ("\nState: ns=" ++ show ns ++ ", sp=" ++  show sp ++ ", x=" ++ show x) r
    where 
          ns = inverseCdf sp x -- Our index in @t@ for the next element
          sp = head $ map snd $ filter ((==) cs . fst) t -- the prob table for cs

inverseCdf :: (Ord a, Num a) => [a] -> a -> Int
inverseCdf pdf x
    | x < head pdf = 0 -- If x is less than the current probability,
                       -- then this is the state we need to return
    | otherwise    = 1+(inverseCdf (tail pdf) (x-(head pdf)))
                     -- Otherwise, subtract x by current probability
                     -- and keep looking

main :: IO ()
main = do
  print $ "Enter a state from " ++ show (map fst table) ++ "> "
  cs <- getLine
  se <- loop (read cs) 5 []
  print se
    where loop _ 0 x = return x
          loop cs c x = do
            n <- randomIO
            let ne = step cs table n
            loop ne (c - 1) $ x ++ [cs]