poem.hs

-- a study for http://morishin.hatenablog.com/entry/haskell-poem-generator
--                                                     2016-03-09 @nfunato

module Lib (generatePoem) where

import System.Random (randomRIO)
import Data.Maybe (maybe)
import Data.List (foldl', tails)
import Data.Map (Map)
import qualified Data.Map as Map
--import qualified Text.MeCab as Mecab

-- main routine

defaultRank = 2

generatePoem :: String -> IO String
generatePoem = generatePoem' defaultRank

generatePoem' :: Int -> String -> IO String
generatePoem' n s = fmap concat . weave . makeTable n =<< parse' s

testMain = putStrLn =<< generatePoem "dummy for parse-dash"

-- parser

type Sentence = [String]
{-
parse  :: String -> IO [Sentence]
parse source = do
  -- assuming an input has one sentence per one line
  mecab <- Mecab.new ["mecab", "-l0"]
  nodeLines <- mapM (Mecab.parseToNodes mecab) (lines source)
  return $ map (filter (not . null) . map Mecab.nodeSurface) nodeLines
-}
parse' :: String -> IO [Sentence]
parse' _ =      -- for test
  return [["サイ", "アーク", "が", "やら", "れ", "た", "よう", "だ", "な", "..."],
          ["フフフ", "...", "奴", "は", "四天王", "中", "でも", "最", "弱", "..."],
          ["人間", "ごとき", "に", "負ける", "と", "は", "魔", "族", "の", "面汚し", "よ", "..."]]

-- prefix table generator

data Phrase = Begin | Middle String | End deriving (Eq, Ord, Show)

undrape :: Phrase -> String
undrape (Middle s) = s
undrape _          = ""

type Prefix = [Phrase]
type PrefixTable = Map Prefix [(Prefix, String)]

makeTable :: Int -> [Sentence] -> PrefixTable
makeTable n = fromList' . concatMap (makeKVs n)

makeKVs :: Int -> Sentence -> [(Prefix, (Prefix, String))]
makeKVs n = map makeKV . slicePhrases (n+1) . toPhrases
  where toPhrases strs = [Begin] ++ map Middle strs ++ [End]
        slicePhrases k = takeWhile ((k==) . length) . map (take k) . tails

makeKV :: [Phrase] -> (Prefix, (Prefix, String))
makeKV  []        = error "makeKV"
makeKV (Begin:xs) = ([Begin], (xs, concatMap undrape xs))
makeKV  xs        = makeKV' xs (last xs)
makeKV' xs End    = ([], ([], "")) -- no essential entry for End
makeKV' xs lst    = (init xs, (tail xs, undrape lst))

fromList' ::(Ord k, Eq a) => [(k, a)] -> Map k [a]
fromList' = foldl' pushnew Map.empty 
  -- simple push might be better than pushnew
  where pushnew mp (k,v) = Map.insertWith f k [v] mp
          where f _ old = if v `elem` old then old else v:old

-- weaver

weave :: PrefixTable -> IO [String]
weave tab = wp [Begin] []
  where wp prefix acc = do
            value <- query prefix tab
            case value of 
                Nothing             -> return $ reverse acc
                Just (next, phrase) -> wp next (phrase:acc)

query :: Prefix -> PrefixTable -> IO (Maybe(Prefix, String))
query prefix tab = maybe (return Nothing) sample (Map.lookup prefix tab) 

sample :: [a] -> IO (Maybe a)
sample [] = return Nothing  -- In fact, [] doesn't appear in this application
sample xs = return . Just . (xs!!) =<< randomRIO (0, length xs - 1)
--sample xs = do idx <- randomRIO (0, length xs - 1); return $ Just (xs !! idx)