module NameGenerator
(Trigrams,
getTrigrams,
first2Chars,
nextChar,
generateName)
where
import qualified Data.Map as Map
import Data.List (isPrefixOf)
import Data.Char (toUpper)
import Text.Read (readMaybe)
import System.Random (randomRIO)
type Trigrams = Map.Map String Int
addTrigram :: Trigrams
-> String
-> Trigrams
addTrigram trigrams line = case (readMaybe value :: Maybe Int) of
Just val -> Map.insert trigram val trigrams
Nothing -> error "There was a problem reading trigrams.txt"
where trigram = take 3 line
value = drop 4 line
getTrigrams :: IO Trigrams
getTrigrams = do
trigrams <- readFile "trigrams.txt"
return $ foldl addTrigram Map.empty $ lines trigrams
filterStartsWith :: String
-> Trigrams
-> Trigrams
filterStartsWith start trigrams = Map.fromList $ filter (\(k, _)->start `isPrefixOf` k) $ Map.toList trigrams
pick :: Int
-> (String, Int)
-> (String, Int)
-> (String, Int)
pick n ("", sumSoFar) (k, v)
| sum' >= n = (k, 0)
| otherwise = ("", sum')
where sum' = sumSoFar + v
pick _ x _ = x
pickFromTrigrams :: Trigrams
-> IO String
pickFromTrigrams trigrams = do
let total = sum $ Map.elems trigrams
chosen <- randomRIO (0, total)
let (trigram, _) = foldl (pick chosen) ("", 0) $ Map.toList trigrams
return trigram
first2Chars :: Trigrams
-> IO String
first2Chars trigrams = do
let options = filterStartsWith " " trigrams
trigram <- pickFromTrigrams options
return $ drop 1 trigram
takeEnd :: Int
-> [a]
-> [a]
takeEnd n l = drop (length l - n) l
nextChar :: Trigrams
-> String
-> IO Char
nextChar trigrams name = do
let options = filterStartsWith (takeEnd 2 name) trigrams
trigram <- pickFromTrigrams options
return $ last trigram
generateName' :: Trigrams
-> IO String
-> IO String
generateName' trigrams currentName = do
name <- currentName
next <- nextChar trigrams name
if next == ' ' then currentName else generateName' trigrams $ fmap (++[next]) currentName
capitalize :: String
-> String
capitalize (first:rest) = (toUpper first):rest
generateName :: Trigrams
-> IO String
generateName trigrams = fmap capitalize $ generateName' trigrams $ first2Chars trigrams