module Game.Mastermind (
   Eval(Eval),
   evaluate,
   matching,
   matchingSimple,

   mixedRandomizedAttempt,
   partitionSizes,

   mainSimple,
   mainRandom,
   main,

   propBestSeparatingCode,
   ) 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
   (intersection, (*&), (#*&), unit, empty, union, unions, cube, )
import Game.Utility (randomSelect, )

import qualified Data.NonEmpty.Set as NonEmptySet
import qualified Data.Map as Map
import qualified Data.Set as Set

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

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

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


data Eval = Eval Int Int
   deriving (Eq, Ord, 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> CodeSet.flatten $ matching (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
{-
   Map.toList $
   Map.mapWithKey
      (\a _ -> Map.update (\n -> toMaybe (n>1) (pred n)) a hist) hist
-}

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

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

-- ToDo: import from combinatorial
{- |
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.
-}
matching :: (CodeSet.C set, Ord a) => Set.Set a -> [a] -> Eval -> set a
matching 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, CodeSetTree.size $ matching 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 -> MS.StateT state Maybe [Char]) ->
   state -> NonEmptySet.T Char -> Int -> IO ()
interaction select initial alphabet n =
   let go state set =
          case MS.runStateT (select set) state of
             Nothing -> putStrLn "contradicting evaluations"
             Just (attempt, newState) -> do
                putStr $ show attempt ++ " " ++
                   show (CodeSet.size set, CodeSet.representationSize set,
                         Set.size (CodeSet.symbols 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 $
                       matching (NonEmptySet.flatten alphabet) attempt $
                       Eval rightPlaces rightSymbols
   in  go initial (cube alphabet n)

mainSimple :: NonEmptySet.T Char -> Int -> IO ()
mainSimple = interaction (MT.lift . listToMaybe . CodeSet.flatten) ()

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

{- |
Remove all but one unused symbols from the alphabet.
-}
reduceAlphabet :: (CodeSet.C set, Ord a) => set a -> Set.Set a -> Set.Set a
reduceAlphabet set alphabet =
   let symbols = CodeSet.symbols set
   in  Set.union symbols $ Set.fromList $ take 1 $ Set.toList $
       Set.difference alphabet symbols

bestSeparatingCode ::
   (CodeSet.C set, Ord a) => Int -> set a -> [[a]] -> [a]
bestSeparatingCode n set =
   let alphabet = CodeSet.symbols set
   in minimax $ \attempt ->
         map (CodeSet.size . intersection set . matching alphabet attempt) $
         possibleEvaluations n

{-
For small sets of codes it is faster to evaluate
all matching codes and build a histogram.
-}
bestSeparatingCodeHistogram ::
   (CodeSet.C set, Ord a) => set a -> [[a]] -> [a]
bestSeparatingCodeHistogram set =
   minimax $ \attempt ->
      Map.elems $ histogram $ map (evaluate attempt) $ CodeSet.flatten set

propBestSeparatingCode ::
   (CodeSet.C set, Ord a) => Int -> set a -> [[a]] -> Bool
propBestSeparatingCode n set attempts =
   bestSeparatingCode n set attempts
   ==
   bestSeparatingCodeHistogram set attempts


{-
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 a -> MS.StateT g Maybe [a]
randomizedAttempt n set = do
   randomAttempts <-
      replicateM 10 $
      replicateM n $
      randomSelect . Set.toList $
      CodeSet.symbols set
   let possible = CodeSet.flatten set
       somePossible =
          -- take 10 possible codes
          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
   _ <- MT.lift $ listToMaybe possible
   return $ bestSeparatingCode n set $ somePossible ++ randomAttempts

{- |
In the beginning we choose codes that separate reasonably well,
based on heuristics.
At the end, when the set becomes small,
we do a brute-force search for an optimally separating code.
-}
{-
The reduced alphabet contains one symbol more than @CodeSet.symbols set@.
Is that necessary or is there always an equally good separating code
without the extra symbol?
-}
separatingRandomizedAttempt ::
   (CodeSet.C set, Rnd.RandomGen g, Ord a) =>
   Int -> Set.Set a -> set a -> MS.StateT g Maybe [a]
separatingRandomizedAttempt n alphabet0 set = do
   case CodeSet.size set of
      0 -> MT.lift Nothing
      1 -> return $ head $ CodeSet.flatten set
      2 -> return $ head $ CodeSet.flatten set
      size ->
         let alphabet = reduceAlphabet set alphabet0
             alphabetSize = Set.size alphabet
             bigSize = toInteger size
         in  if bigSize * (bigSize + toInteger alphabetSize ^ n) <= 1000000
               then return $ bestSeparatingCodeHistogram set $
                    CodeSet.flatten set ++ replicateM n (Set.toList alphabet)
               else randomizedAttempt n set

{- |
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 a -> MS.StateT g Maybe [a]
mixedRandomizedAttempt n set = do
   case CodeSet.size set of
      0 -> MT.lift Nothing
      1 -> return $ head $ CodeSet.flatten set
      2 -> return $ head $ CodeSet.flatten set
      size ->
         if size <= 100
           then randomizedAttempt n set
           else
              fmap (CodeSet.select set) $
              MS.state $ Rnd.randomR (0, size-1)

mainRandom :: NonEmptySet.T Char -> Int -> IO ()
mainRandom alphabet n = do
   g <- Rnd.getStdGen
   interaction
      (separatingRandomizedAttempt n (NonEmptySet.flatten alphabet))
      g alphabet n

main :: IO ()
main =
   let alphabet = NonEmptySet.fromList ('a'!:['b'..'z'])
   in  if True
         then mainRandom alphabet 5
         else mainSimple alphabet 7

{-
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" $ matching (Set.fromList $ ['a'..'z']) attempt (evaluate "amiga" attempt)) ["uvqcm","wukjv","lmoci","caoab","mbadi","ombed","lqbia"]
[True,True,True,True,False,True,False]
-}