{-# LANGUAGE PatternGuards, ViewPatterns #-}

module Solve where

import Data.Array
import Data.Tree
import Data.Function (on)
import Data.List (groupBy, maximumBy)
import Data.Maybe (isNothing, isJust)
import Data.Monoid
import Control.Applicative (liftA2)
import Control.Arrow ((&&&))
import Control.Monad (guard)

data Player = X | O
  deriving (Show, Eq, Ord)

next X = O
next O = X

data Result = Win Player Int [Loc] -- ^ This player can win in n moves
            | Cats                 -- ^ Tie game
  deriving (Show, Eq)

compareResultsFor :: Player -> (Result -> Result -> Ordering)
compareResultsFor X = compare `on` resultToScore
    where resultToScore (Win X n _) = (1/(1+fromIntegral n))
          resultToScore Cats        = 0
          resultToScore (Win O n _) = (-1/(1+fromIntegral n))
compareResultsFor O = flip (compareResultsFor X)

type Loc = (Int,Int)
type Board = Array Loc (Maybe Player)

emptyBoard :: Board
emptyBoard = listArray ((0,0), (2,2)) (repeat Nothing)

showBoard :: Board -> String
showBoard = unlines . map showRow . groupBy ((==) `on` (fst . fst)) . assocs
  where showRow = concatMap (showPiece . snd)
        showPiece Nothing  = " "
        showPiece (Just p) = show p

data Move = Move Player Loc
  deriving Show

makeMove :: Move -> Board -> Board
makeMove (Move p l) b = b // [(l, Just p)]

data Game = Game Board           -- ^ The current board state.
                 Player          -- ^ Whose turn is it?
                 [Move]          -- ^ The list of moves so far (most
                                 --   recent first).
  deriving Show

initialGame = Game emptyBoard X []

-- | The full game tree for tic-tac-toe.
gameTree :: Tree Game
gameTree = unfoldTree (id &&& genMoves) initialGame

-- | Generate all possible successor games from the given game.
genMoves :: Game -> [Game]
genMoves (Game board player moves) = newGames
  where validLocs = map fst . filter (isNothing . snd) . assocs $ board
        newGames  = [Game (makeMove m board) (next player) (m:moves)
                      | p <- validLocs
                      , let m = Move player p
                    ]

-- | Simple fold for Trees.  The Data.Tree module does not provide
--   this.
foldTree :: (a -> [b] -> b) -> Tree a -> b
foldTree f (Node a ts) = f a (map (foldTree f) ts)

-- | Solve the game for player @p@: prune all but the optimal moves
--   for player @p@, and annotate each game with its result (given
--   best play).
solveFor :: Player -> Tree Game -> Tree (Game, Result)
solveFor p = foldTree (solveStep p)

-- | Given a game and its continuations (including their results),
--   solve the game for player p.  If it is player p's turn, prune all
--   continuations except the optimal one for p. Otherwise, leave all
--   continuations.  The result of this game is the result of the
--   optimal choice if it is p's turn, otherwise the worst possible
--   outcome for p.
solveStep :: Player -> Game -> [Tree (Game, Result)] -> Tree (Game, Result)
solveStep p g@(Game brd curPlayer moves) conts
  | Just res <- gameOver g = Node (g, res) []

  | curPlayer == p = let c   = bestContFor p conts
                         res = inc . snd . rootLabel $ c
                     in  Node (g, res) [c]
  | otherwise      = Node (g, bestResultFor (next p) conts) conts

bestContFor :: Player -> [Tree (Game, Result)] -> Tree (Game, Result)
bestContFor p = maximumBy (compareResultsFor p `on` (snd . rootLabel))

bestResultFor :: Player -> [Tree (Game, Result)] -> Result
bestResultFor p = inc . snd . rootLabel . bestContFor p

inc :: Result -> Result
inc (Win p n ls) = Win p (n+1) ls
inc Cats         = Cats

-- | Check whether the game is over, returning the result if it is.
gameOver :: Game -> Maybe Result
gameOver (Game board _ _)
  = getFirst $ mconcat (map (checkWin board) threes) `mappend` checkCats board

checkWin :: Board -> [Loc] -> First Result
checkWin board = First
               . (>>= winAsResult)      -- Maybe Result
               . mapM strength          -- Maybe [(Loc, Player)]
               . map (id &&& (board!))  -- [(Loc, Maybe Player)]

winAsResult :: [(Loc, Player)] -> Maybe Result
winAsResult (unzip -> (ls,ps))
  | Just p <- allEqual ps = Just (Win p 0 ls)
winAsResult _ = Nothing

checkCats :: Board -> First Result
checkCats b | all isJust (elems b) = First (Just Cats)
            | otherwise            = First Nothing

allEqual :: Eq a => [a] -> Maybe a
allEqual = foldr1 combine . map Just
  where combine (Just x) (Just y) | x == y = Just x
                                  | otherwise = Nothing
        combine Nothing _         = Nothing
        combine _ Nothing         = Nothing

strength :: Functor f => (a, f b) -> f (a,b)
strength (a, f) = fmap ((,) a) f

threes :: [[Loc]]
threes = rows ++ cols ++ diags
  where rows     = [ [ (r,c) | c <- [0..2] ] | r <- [0..2] ]
        cols     = [ [ (r,c) | r <- [0..2] ] | c <- [0..2] ]
        diags    = [ [ (i,i) | i <- [0..2] ]
                   , [ (i,2-i) | i <- [0..2] ]
                   ]