{-
The LetterGuesser module is responsible for the AI process of selecting
the next letter to be guessed each round of Hangman. LetterGuesser uses
two pruning strategies depending on whether a guessed letter was in the
word being guessed or not.
-}

module LetterGuesser where

import           Control.Monad(liftM2)
import           Types
import           Data.Map (Map, fromList, insertWith, lookup)
import           Data.Set (Set, member, difference, fromList, toList)

-- Update dictionary of words if this correct letter was guessed
--   by removing words that don't match the updated letter positions
pruneCorrect :: InternalState -> [Types.Word]
pruneCorrect is = filter positionAgreement dictWords
    where dictWords = wordDict is
          validLetter maybeLetter dictWordLetter = case maybeLetter of
              (Just l) -> dictWordLetter == l
              _ -> True
          positionAgreement w = and $ zipWith validLetter (currentWord is) (getWordString w)


-- Update dictionary of words if this incorrect letter was guessed
--   by removing words that have the incorrect letter
pruneIncorrect :: InternalState -> Char -> [Types.Word]
pruneIncorrect is c = filter wordsWithLetter possibleWords
    where possibleWords = wordDict is
          wordsWithLetter w = not $ elem c (getWordString w)

-- Top-level function that uses AI strategies to determine the 
--   best letter to guess given the game's InteralState
selectLetter :: InternalState -> Char
selectLetter is = mostCommonLetter letterFrequencies unguessedLetters unguessedLettersList (head unguessedLettersList)
    where possibleWords = wordDict is
          guessedLetters = usedLetters is
          unguessedLetters = difference (Data.Set.fromList alphabet) guessedLetters
          unguessedLettersList = toList $ difference (Data.Set.fromList alphabet) guessedLetters
          letterFrequencies = foldr (updateLetterCountsWithWord guessedLetters) emptyLetterMap possibleWords

-- Selects the most common letter from the given letter frequency map
mostCommonLetter :: Map Char Int -> Set Char -> [Char] -> Char -> Char
mostCommonLetter _ _ [] mostFreqLetter = mostFreqLetter
mostCommonLetter letterFreqs unguessedLetters (currentLetter:xs) mostFreqLetter = 
    case liftM2 (>) (freqCount mostFreqLetter) (freqCount currentLetter) of
      Just (True) -> mostCommonLetter letterFreqs unguessedLetters xs mostFreqLetter
      _ -> mostCommonLetter letterFreqs unguessedLetters xs currentLetter
    where freqCount l = Data.Map.lookup l letterFreqs

-- Update the letter frequency map with the letters in the given word
updateLetterCountsWithWord :: Set Char -> Types.Word -> Map Char Int -> Map Char Int
updateLetterCountsWithWord guessedLetters (Types.Word w) m = foldr (incLetterCount guessedLetters) m w

-- Increment the frequency of the given letter in the letter frequency map
incLetterCount :: Set Char -> Char -> Map Char Int -> Map Char Int
incLetterCount guessedLetters c m
    | member c guessedLetters = m
    | otherwise = insertWith (+) c 1 m

emptyLetterMap :: Map Char Int
emptyLetterMap = Data.Map.fromList $ map (\l -> (l, 0)) alphabet