module Markov where

import Control.Monad (guard)
import Data.Monoid
import System.Random (StdGen,mkStdGen,getStdGen,setStdGen,randomR,randomRIO)
import System.IO.Unsafe (unsafeInterleaveIO)
import qualified Data.Foldable as F
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Text.Lazy as L
import qualified Data.Text.Lazy.IO as L

import Debug.Trace


data Markov m a = Markov
  { mStart :: Set.Set a
    -- ^ Set of starting atoms
  , mNext  :: Map.Map a (Map.Map a m)
    -- ^ Map of transitions between these two atoms
  } deriving (Show)

instance (Monoid m, Ord a) => Monoid (Markov m a) where
  mempty  = Markov
    { mStart = Set.empty
    , mNext  = Map.empty
    }

  mappend l r = Markov
    { mStart = Set.union (mStart l) (mStart r)
    , mNext  = Map.unionWith combine (mNext l) (mNext r)
    }

combine :: (Monoid m, Ord a) => Map.Map a m -> Map.Map a m -> Map.Map a m
combine  = Map.unionWith mappend


-- Learning Mode ---------------------------------------------------------------

type MarkovL = Markov (Sum Int)

fromList :: Ord a => [a] -> MarkovL a
fromList as0 = case as0 of
  a:rest -> loop (start a mempty) a rest
  _      -> error "fromList: empty input"
  where
  loop m a as = case as of
    b:rest -> loop (next a b m) b rest
    []     -> m

-- | Add an atom to the starting set.
start :: Ord a => a -> MarkovL a -> MarkovL a
start a m = m { mStart = Set.insert a (mStart m) }

-- | Increment the edge count between two atoms.
next :: Ord a => a -> a -> MarkovL a -> MarkovL a
next a a' m = m
  { mNext = Map.insertWith combine a (Map.singleton a' (Sum 1)) (mNext m)
  }


-- Compiled Mode ---------------------------------------------------------------

type MarkovC = Markov Double

compile :: Ord a => MarkovL a -> MarkovC a
compile m = m { mNext = Map.mapWithKey update (mNext m) }
  where
  update k as = Map.map (\p -> fromIntegral (getSum p) / total) as
    where
    total = fromIntegral (getSum (F.fold (Map.elems as)))

stepsIO :: Ord a => MarkovC a -> a -> IO [a]
stepsIO s a0 = (a0:) `fmap` loop a0
  where
  loop a = do
    mb <- stepIO s a
    case mb of
      Just a' -> do
        rest <- unsafeInterleaveIO (loop a')
        return (a':rest)
      Nothing -> return []

stepsIO' :: Ord a => MarkovC a -> IO [a]
stepsIO' s = do
  let starts = Set.toList (mStart s)
  ix <- randomRIO (0,length starts - 1)
  stepsIO s (starts !! ix)

stepIO :: Ord a => MarkovC a -> a -> IO (Maybe a)
stepIO s a = do
  r <- getStdGen
  case step s a r of
    Just (a',r') -> do
      setStdGen r'
      return (Just a')
    Nothing -> return Nothing

step :: Ord a => MarkovC a -> a -> StdGen -> Maybe (a,StdGen)
step s a r = do
  ns <- Map.lookup a (mNext s)
  choose (Map.toList ns) r

choose :: [(a,Double)] -> StdGen -> Maybe (a,StdGen)
choose as r = do
  guard (not (null as))
  let (val,r') = randomR (0,1) r
      loop p xs = case xs of
        (a,ap):rest | ap > p    -> Just (a,r')
                    | otherwise -> loop (p - ap) rest
        []                      -> Nothing
  loop val as



-- Tests -----------------------------------------------------------------------

parse :: FilePath -> IO (MarkovC L.Text)
parse path = do
  input <- L.readFile path
  return $ compile
         $ mconcat
         $ map (fromList . L.words)
         $ filter (not . L.null)
         $ L.lines input

sample :: MarkovC L.Text -> IO ()
sample m = do
  ws <- stepsIO' m
  L.putStrLn (L.unwords ws)