A simple Markov-Chain implementation

markov2.hs

{-# LANGUAGE TupleSections #-}

import Data.List
import Control.Monad.Random
import Control.Arrow
import Control.Applicative
import Safe
import System.Environment
import System.Exit
import Data.Hashable (Hashable)
import qualified Data.HashMap.Strict as H

type Lookup a b = H.HashMap a b
type MC a = Lookup [a] (Lookup a Int)

main :: IO ()
main = getArgs >>= run . map readMay

help :: IO ()
help = putStrLn "Usage: markov CONTEXT_LENGTH" >> exitFailure

run :: [Maybe Int] -> IO ()
run [Just n] = do
  t <- getContents
  evalRandIO (runChain n t) >>= putStrLn

run _ = help

window :: Int -> [a] -> [[a]]
window n = map (take n) . transpose . tails

maperizeList :: (Ord a, Hashable a) => Int -> [a] -> ([a], Lookup a Int)
maperizeList n l = (take n l, H.singleton (last l) 1)

buildMC :: (Ord a, Hashable a) => Int -> [[a]] -> MC a
buildMC n = H.fromListWith (H.unionWith (+)) . map (maperizeList n)

buildMCFromList :: (Ord a, Hashable a) => Int -> [a] -> MC a
buildMCFromList n = buildMC n . window (succ n)

chainItem :: (Ord a, Hashable a, RandomGen g) => MC a -> [a] -> Rand g [a]
chainItem m l = case H.lookup l m of
                     Nothing -> randomKey m >>= chainItem m
                     Just t  -> do i <- randomWeightedItem t
                                   return $ drop 1 l ++ [i]

iterateM :: (Functor m, Monad m) => (a -> m a) -> a -> m [a]
-- iterateM f = sequence . drop 1 . iterate (f =<<) . return -- Isn't efficient for some reason
iterateM f a = do x <- f a
                  (x:) <$> iterateM f x

chainItems :: (RandomGen g, Ord a, Hashable a) => MC a -> [a] -> Rand g [[a]]
chainItems m = iterateM (chainItem m)

runChain :: (RandomGen g, Ord a, Hashable a) => Int -> [a] -> Rand g [a]
runChain n l = map head <$> chainItems (buildMCFromList n l) []

randomWeightedItem :: RandomGen g => Lookup a Int -> Rand g a
randomWeightedItem m = fromList $ map (second realToFrac) (H.toList m)

randomKey :: RandomGen g => Lookup a b -> Rand g a
randomKey m = uniform (H.keys m)