module Game.Mastermind where

import qualified Game.Mastermind.CodeSet.Tree as CodeSetTree
-- import qualified Game.Mastermind.CodeSet.Union as CodeSetUnion
import qualified Game.Mastermind.CodeSet as CodeSet
import Game.Mastermind.CodeSet
   (flatten, intersection, (*&), (#*&), unit, empty, union, unions, cube, )

import qualified Data.Map as Map
import qualified Data.Set as Set

import Data.List.HT (partition, )
import Data.Tuple.HT (mapPair, )
import Data.Maybe.HT (toMaybe, )
import Data.Maybe (listToMaybe, )
import Control.Monad (liftM, guard, when, replicateM, )

import qualified Control.Monad.Trans.State as State
import qualified Control.Monad.Trans.Class as Trans

import qualified System.Random as Rnd
import qualified System.IO as IO


data Eval = Eval Int Int
   deriving (Eq, Show)

{- |
Given the code and a guess, compute the evaluation.
-}
evaluate :: (Ord a) => [a] -> [a] -> Eval
evaluate code attempt =
   uncurry Eval $
   mapPair
      (length,
       sum . Map.elems .
       uncurry (Map.intersectionWith min) .
       mapPair (histogram,histogram) . unzip) $
   partition (uncurry (==)) $
   zip code attempt

{-
*Game.Mastermind> filter ((Eval 2 0 ==) . evaluate "aabbb") $ replicateM 5 ['a'..'c']
["aaaaa","aaaac","aaaca","aaacc","aacaa","aacac","aacca","aaccc","acbcc","accbc","acccb","cabcc","cacbc","caccb","ccbbc","ccbcb","cccbb"]
*Game.Mastermind> flatten $ remaining (Set.fromList ['a'..'c']) "aabbb" (Eval 2 0)
["aaaaa","aaaac","aaaca","aaacc","aacaa","aacac","aacca","aaccc","acbcc","accbc","acccb","cabcc","cacbc","caccb","ccbbc","ccbcb","cccbb"]
-}


histogram :: (Ord a) => [a] -> Map.Map a Int
histogram =
   Map.fromListWith (+) . map (\a -> (a,1))

selectFromHistogram :: (Ord a) => Map.Map a Int -> [(a, Map.Map a Int)]
selectFromHistogram hist =
   map (\a -> (a, Map.update (\n -> toMaybe (n>1) (pred n)) a hist)) $
   Map.keys hist

{- |
A variant of the game:
It is only possible to specify number of symbols at right places.

The results of 'remaining' and 'remainingSimple' cannot be compared.
-}
remainingSimple :: Ord a => Set.Set a -> [a] -> Int -> [[Set.Set a]]
remainingSimple alphabet code rightPlaces =
   map
      (zipWith
         (\symbol right ->
            if right
              then Set.singleton symbol
              else Set.delete symbol alphabet)
         code) $
   possibleRightPlaces (length code) rightPlaces

{- |
Combinatorical \"choose k from n\".
-}
possibleRightPlaces :: Int -> Int -> [[Bool]]
possibleRightPlaces n rightPlaces =
   if n < rightPlaces
     then []
     else
       if n==0
         then [[]]
         else
            (guard (rightPlaces>0) >>
               (map (True:) $
                possibleRightPlaces (n-1) (rightPlaces-1)))
            ++
            (map (False:) $
             possibleRightPlaces (n-1) rightPlaces)

{- |
Given a code and an according evaluation,
compute the set of possible codes.

The Game.Mastermind game consists of collecting pairs
of codes and their evaluations.
The searched code is in the intersection of all corresponding code sets.
-}
remaining ::
   (CodeSet.C set, Ord a) =>
   Set.Set a -> [a] -> Eval -> set a
remaining alphabet =
   let findCodes =
          foldr
             (\(fixed,c) go rightSymbols floating0 ->
                if fixed
                  then c #*& go rightSymbols floating0
                  else
                    (unions $ do
                        guard (rightSymbols > 0)
                        (src, floating1) <- selectFromHistogram floating0
                        guard (c /= src)
                        return $ src #*& go (rightSymbols-1) floating1)
                    `union`
                    (Set.difference
                        (Set.delete c alphabet)
                        (Map.keysSet floating0) *&
                     go rightSymbols floating0))
             (\rightSymbols _floating ->
                if rightSymbols>0
                  then empty
                  else unit)
   in \code (Eval rightPlaces rightSymbols) ->
       unions $
       map
          (\pattern ->
             let patternCode = zip pattern code
             in  findCodes patternCode rightSymbols $
                 histogram $ map snd $ filter (not . fst) patternCode) $
       possibleRightPlaces (length code) rightPlaces


partitionSizes :: (Ord a) => Set.Set a -> [a] -> [(Eval, Integer)]
partitionSizes alphabet code =
   map (\eval ->
      (eval,
       CodeSet.size $
       (id :: CodeSetTree.T a -> CodeSetTree.T a) $
       remaining alphabet code eval)) $
   possibleEvaluations (length code)

possibleEvaluations :: Int -> [Eval]
possibleEvaluations n = do
   rightPlaces <- [0..n]
   rightSymbols <- [0..n-rightPlaces]
   return $ Eval rightPlaces rightSymbols


interaction ::
   (CodeSetTree.T Char -> State.StateT state Maybe [Char]) ->
   state ->
   Set.Set Char -> Int -> IO ()
interaction select initial alphabet n =
   let go state set =
          case State.runStateT (select set) state of
             Nothing -> putStrLn "contradicting evaluations"
             Just (attempt, newState) -> do
                putStr $ show attempt ++ " " ++
                   show (CodeSet.size set, CodeSet.representationSize set) ++ " "
                IO.hFlush IO.stdout
                eval <- getLine
                let evalHist = histogram eval
                    evalHistRem =
                       Map.keys $ Map.delete 'o' $ Map.delete 'x' evalHist
                when (not $ null evalHistRem)
                   (putStrLn $ "ignoring: " ++ evalHistRem)
                let rightPlaces  = length (filter ('x' ==) eval)
                    rightSymbols = length (filter ('o' ==) eval)
                if rightPlaces >= n
                  then putStrLn "I won!"
                  else go newState $ intersection set $
                       remaining alphabet attempt $
                       Eval rightPlaces rightSymbols
   in  go initial (cube alphabet n)

mainSimple :: Set.Set Char -> Int -> IO ()
mainSimple =
   interaction
      (Trans.lift . listToMaybe . flatten)
      ()

-- candidate for random-utility, cf. module htam:Election, markov-chain
randomSelect ::
   (Rnd.RandomGen g, Monad m) =>
   [a] -> State.StateT g m a
randomSelect items =
   liftM (items!!) $ State.StateT $ return . Rnd.randomR (0, length items-1)

{- |
minimum of maximums using alpha-beta-pruning
-}
minimax :: (Ord b) => [(a, [b])] -> a
minimax [] = error "minimax of empty list"
minimax ((a0,bs0):rest) =
   fst $
   foldl
      (\old@(_minA, minB) (a,bs) ->
         let (ltMinB, gtMinB) = partition (minB>) bs
         in  if null gtMinB
               then (a, maximum ltMinB)
               else old)
      (a0, maximum bs0) rest

{-
Here we optimize for small set sizes.
For performance we could optimize for small set representation sizes.
However the resulting strategy looks much like the strategy
from mainSimple and needs more attempts.
-}
randomizedAttempt ::
   (CodeSet.C set, Rnd.RandomGen g, Ord a) =>
   Int -> Set.Set a ->
   set a -> State.StateT g Maybe [a]
randomizedAttempt n alphabet set = do
   randomAttempts <-
      replicateM 10 $
      replicateM n $
      randomSelect . Set.toList $
      CodeSet.symbols set
   let possible = flatten set
       somePossible =
          -- take 10 possible
          let size = CodeSet.size set
              num = 10
          in  map (CodeSet.select set) $
              Set.toList $ Set.fromList $
              take num $
              map (flip div (fromIntegral num)) $
              iterate (size+) 0
   _ <- Trans.lift $ listToMaybe possible
   return $ minimax $
      map
         (\attempt -> (attempt,
            map (CodeSet.size . intersection set . remaining alphabet attempt) $
            possibleEvaluations n)) $
      somePossible ++ randomAttempts

{- |
In the beginning we simply choose a random code
from the set of possible codes.
In the end, when the set becomes small,
then we compare different alternatives.
-}
mixedRandomizedAttempt ::
   (CodeSet.C set, Rnd.RandomGen g, Ord a) =>
   Int -> Set.Set a ->
   set a -> State.StateT g Maybe [a]
mixedRandomizedAttempt n alphabet set = do
   case CodeSet.size set of
      0 -> Trans.lift Nothing
      1 -> return $ head $ CodeSet.flatten set
      2 -> return $ head $ CodeSet.flatten set
      size ->
         if size <= 100
           then randomizedAttempt n alphabet set
           else
              fmap (CodeSet.select set) $
              State.StateT $ return . Rnd.randomR (0, size-1)

mainRandom :: Set.Set Char -> Int -> IO ()
mainRandom alphabet n = do
   g <- Rnd.getStdGen
   interaction (randomizedAttempt n alphabet) g alphabet n

main :: IO ()
main =
--   mainSimple (Set.fromList ['a'..'z']) 7
   mainRandom (Set.fromList ['a'..'z']) 5

{-
Bug: (fixed)
*Game.Mastermind> main
"uvqcm" (11881376,130) o
"wukjv" (3889620,440)
"lmoci" (1259712,372) xo
"caoab" (94275,1765) oo
"mbadi" (6856,2091) ooo
"ombed" (327,447) x
"lqbia" (2,10) xo
contradicting evaluations
*Game.Mastermind> map (evaluate "amiga") ["uvqcm","wukjv","lmoci","caoab","mbadi","ombed","lqbia"]
[Eval 0 1,Eval 0 0,Eval 1 1,Eval 0 2,Eval 0 3,Eval 1 0,Eval 1 1]
*Game.Mastermind> map (\attempt -> member "amiga" $ remaining (Set.fromList $ ['a'..'z']) attempt (evaluate "amiga" attempt)) ["uvqcm","wukjv","lmoci","caoab","mbadi","ombed","lqbia"]
[True,True,True,True,False,True,False]
-}