{- |
https://en.wikipedia.org/wiki/Mastermind_(board_game)

Given a list of guesses and according evaluations,
the solver computes a list of all possible codes
that match the obtained evaluations.

See also the @board-games@ package.
-}
module Main where

import qualified Math.SetCover.Exact as ESC

import qualified System.IO as IO
import System.Random (StdGen, getStdGen, randomR, )

import qualified Control.Monad.Trans.State as MS
import Control.Monad (liftM2, replicateM, when, )

import qualified Data.Set as Set; import Data.Set (Set, )
import qualified Data.Array as Array
import qualified Data.List.Match as Match
import qualified Data.List.HT as ListHT
import Data.Tuple.HT (mapSnd, )
import Data.List.HT (tails, viewL, viewR, )
import Data.Maybe (mapMaybe, )


-- cf. htam:Combinatorics.tuples
choose :: Int -> [a] -> [[a]]
choose n xs =
   flip MS.evalStateT xs $ replicateM n $
   MS.StateT $ mapMaybe viewL . tails


data X = Pos Int | Eval Eval Int Int | EvalRow Eval Int
        deriving (Eq, Ord, Show)

data Eval = CorrectPlace | CorrectSymbol
        deriving (Eq, Ord, Show)

type Assign a = ESC.Assign [(Int, a)] (Set X)

assignsFromGuesses ::
   (Ord a) =>
   Int -> [a] -> [([a], (Int,Int))] -> [Assign a]
assignsFromGuesses width set guesses =
   liftM2
      (\pat a ->
         let ks = map fst $ filter snd $ zip [0..] pat
         in  ESC.assign (map (flip (,) a) ks) $ Set.unions $
             Set.fromList (map Pos ks) :
             zipWith
                (\row (guess,_) ->
                   Set.fromList $
                   let (correctlyPlaced, remGuess) =
                          ListHT.partition (\(_k, (used,equ)) -> used && equ) $
                          zip [0..] $ zip pat $ map (a==) guess
                   in  map (Eval CorrectPlace row . fst) correctlyPlaced
                       ++
                       map (Eval CorrectSymbol row . fst)
                          (Match.take
                             (filter (fst . snd) remGuess)
                             (filter (snd . snd) remGuess)))
                [0..] guesses)
      (tail $ replicateM width [False, True]) set
   ++
   concat
      (zipWith
         (\row (_, (correctPlaces,correctSymbols)) ->
            let fill eval k =
                   map (ESC.assign [] . Set.fromList . (EvalRow eval row :)) $
                   choose (width - k) $
                   map (Eval eval row) $ take width [0..]
            in  fill CorrectPlace correctPlaces
                ++
                fill CorrectSymbol correctSymbols)
         [0..] guesses)


codeFromLabels :: [[(Int, a)]] -> [a]
codeFromLabels mxs =
   case concat mxs of
      xs -> Array.elems $ Array.array (0, length xs - 1) xs


unique :: (Ord a) => [a] -> Bool
unique xs = Set.size (Set.fromList xs) == length xs

newGuess ::
   (Ord a) =>
   Int -> [a] -> [([a], (Int,Int))] -> MS.State StdGen (Maybe [a])
newGuess width alphabet oldGuesses = do
   n <- MS.state $ randomR (1,1000)
   return $ fmap snd $ viewR $ take n $
--      filter unique $
      map codeFromLabels $ ESC.partitions $
      assignsFromGuesses width alphabet oldGuesses

countEval :: MS.State String (Int, Int)
countEval =
   let count c = fmap length $ MS.state $ ListHT.partition (c==)
   in  liftM2 (,) (count 'x') (count 'o')

{- |
In every round the computer player selects randomly one of the first 1000 codes
that are coherent with the known evaluations.
This strategy prevents stupid guesses like "aaaaa",
but it does not minimize the number of guesses.
When the game approaches the end
there is often only one unknown letter left
and the algorithm makes a guess for ruling out every single candidate.
It would be more efficient to use non-coherent guesses in this situation
in order to rule out a whole bunch of candidates at once.
-}
interaction :: Int -> [Char] -> IO ()
interaction width alphabet =
   let go guesses g0 =
          case MS.runState (newGuess width alphabet guesses) g0 of
             (Nothing, _) -> putStrLn "contradicting evaluations"
             (Just attempt, g1) -> do
                putStr $ show attempt ++ " "
                IO.hFlush IO.stdout
                eval0 <- getLine
                let ((numPlaces, numSymbols), evalRem) =
                      MS.runState countEval eval0
                when (not $ null evalRem) (putStrLn $ "ignoring: " ++ evalRem)
                if numPlaces >= width
                  then putStrLn "Code found!"
                  else go ((attempt, (numPlaces, numSymbols)) : guesses) g1
   in  go [] =<< getStdGen

testGuesses :: [(String, (Int, Int))]
testGuesses =
   map (mapSnd (MS.evalState countEval)) $
   ("aaaayw", "x") :
   ("bbbdcw", "") :
   ("eefeym", "oo") :
   ("iuzamf", "oo") :
   ("gvarfe", "ooo") :
   ("paqfes", "xxo") :
   ("vamsej", "ooxx") :
   ("amgses", "ooox") :
   ("majgep", "xxx") :
   []

testSolve :: IO ()
testSolve =
   mapM_ (print . codeFromLabels) $ ESC.partitions $
   assignsFromGuesses 6 ['a'..'z'] testGuesses


main :: IO ()
main = do
   let n = 5
   putStrLn $
      "Come up with a word consisting of " ++ show n ++
      " letters and evaluate my guesses."
   putStrLn "Enter 'x's for correct places and 'o's for correct symbols in any order."
   interaction n ['a'..'z']