{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

import Control.Arrow
import Control.Monad
import Data.Array
import Data.Function
import Data.List
import Data.Maybe
import System.Environment

(|>) :: a -> (a -> b) -> b
(|>) = flip ($)

mapWhile :: (a -> Maybe b) -> [a] -> [b]
mapWhile _ [] = []
mapWhile f (x:xs) =
    case f x of
        Just y  -> y : mapWhile f xs
        Nothing -> []

compose :: [a -> a] -> a -> a
compose = flip (foldl (flip id))

maximumOn :: Ord b => (a -> b) -> [a] -> Maybe a
maximumOn _ [] = Nothing
maximumOn f xs = Just $ maximumBy (compare `on` f) xs

minimumOn :: Ord b => (a -> b) -> [a] -> Maybe a
minimumOn _ [] = Nothing
minimumOn f xs = Just $ minimumBy (compare `on` f) xs

data Color
    = Black
    | White
    deriving (Show, Eq)

other :: Color -> Color
other Black = White
other White = Black

prettyColor :: Color -> String
prettyColor Black = "X"
prettyColor White = "O"

prettyTile :: Maybe Color -> String
prettyTile = maybe "." prettyColor

data Board
    = Board (Array ColRow (Maybe Color))
    deriving Eq

instance Show Board where
    show = prettyBoard

prettyBoard :: Board -> String
prettyBoard (Board b) = concatMap prettyRow is
    where
        (_, (n, m)) = bounds b
        is = range (1, n)
        js = range (1, m)
        prettyRow i   = unwords (map (prettyCol i) js) ++ "\n\n"
        prettyCol i j = prettyTile (b ! (i, j))

type ColRow = (Int, Int)

type BoardSize = (Int, Int)

boardOf :: BoardSize -> [Maybe Color] -> Board
boardOf size = Board . listArray ((1, 1), size)

emptyBoard :: BoardSize -> Board
emptyBoard size = boardOf size (repeat Nothing)

isEmpty :: Board -> ColRow -> Bool
isEmpty (Board b) ij = isNothing (b ! ij)

standardBoard :: Board
standardBoard =
    emptyBoard (8, 8)
    `placing` Piece (4, 4) White
    `placing` Piece (4, 5) Black
    `placing` Piece (5, 4) Black
    `placing` Piece (5, 5) White

score :: Board -> (Int, Int)
score (Board b) = foldl count (0, 0) b
    where
        count (x, y) (Just Black) = (x + 1, y)
        count (x, y) (Just White) = (x, y + 1)
        count (x, y) _            = (x, y)

relScoreFor :: Color -> Board -> Int
relScoreFor Black board = let (x, y) = score board in x - y
relScoreFor White board = let (x, y) = score board in y - x

data Piece
    = Piece ColRow Color
    deriving (Show, Eq)

place :: Piece -> Board -> Board
place (Piece ij color) (Board b) = Board (b // [(ij, Just color)])

placing = flip place

swap :: Piece -> Piece
swap (Piece ij color) = Piece ij (other color)

capture :: [Piece] -> Board -> Board
capture pieces = compose $ fmap (place . swap) pieces

isColor :: Piece -> Color -> Bool
Piece _ x `isColor` y = x == y

at :: Board -> ColRow -> Maybe (Maybe Piece)
Board b `at` ij
    | inRange (bounds b) ij = Just $ fmap (Piece ij) (b ! ij)
    | otherwise             = Nothing

data Direction
    = DirN
    | DirNE
    | DirE
    | DirSE
    | DirS
    | DirSW
    | DirW
    | DirNW
    deriving (Show, Eq, Ord, Enum)

directions :: [Direction]
directions = enumFrom (toEnum 0)

stepTowards :: Direction -> ColRow -> ColRow
stepTowards DirN  (i, j) = (i - 1, j)
stepTowards DirNE (i, j) = (i - 1, j + 1)
stepTowards DirE  (i, j) = (i, j + 1)
stepTowards DirSE (i, j) = (i + 1, j + 1)
stepTowards DirS  (i, j) = (i + 1, j)
stepTowards DirSW (i, j) = (i + 1, j - 1)
stepTowards DirW  (i, j) = (i, j - 1)
stepTowards DirNW (i, j) = (i - 1, j - 1)

walkTowards :: Direction -> ColRow -> [ColRow]
walkTowards dir start = tail (iterate (stepTowards dir) start)

longestCaptureFor :: Color -> [Piece] -> [Piece]
longestCaptureFor me pieces =
    let (capture, rest) = break (`isColor` me) pieces
        in if null rest then [] else capture

untilBounds :: Board -> [ColRow] -> [Maybe Piece]
untilBounds board = mapWhile (board `at`)

takeString :: [Maybe Piece] -> [Piece]
takeString = mapWhile id

capturesAt :: Board -> Color -> ColRow -> [Piece]
capturesAt board me (i, j) = do
    dir <- directions
    walkTowards dir (i, j)
    |> untilBounds board
    |> takeString
    |> longestCaptureFor me

capturesAt' board me ij = (ij, capturesAt board me ij)

positions :: Board -> [ColRow]
positions (Board b) = indices b

freePositions :: Board -> [ColRow]
freePositions board = filter (isEmpty board) (positions board)

validPositions :: Color -> Board -> [(ColRow, [Piece])]
validPositions me board =
    freePositions board
    |> fmap (capturesAt' board me)
    |> filter (not . null . snd)

validMoves :: Color -> Board -> [Board]
validMoves me board =
    validPositions me board
    |> fmap (\(ij, ps) -> capture ps board `placing` Piece ij me)

data Reversi
    = Reversi
    { reversiPlayer :: Color
    , reversiBoard  :: Board
    }

reversiInit :: Reversi
reversiInit = Reversi Black standardBoard

reversiScore = score . reversiBoard

reversiValidPositions rev = validPositions (reversiPlayer rev) (reversiBoard rev)

reversiMove :: ColRow -> Reversi -> Reversi
reversiMove ij rev =
    Reversi (other me)
            (capture (capturesAt board me ij) board `placing` Piece ij me)
    where
        me    = reversiPlayer rev
        board = reversiBoard rev

reversiNullMove :: Reversi -> Reversi
reversiNullMove rev = Reversi (other $ reversiPlayer rev) (reversiBoard rev)

instance Show Reversi where
    show = prettyBoard . reversiBoard

class Game g where
    type AgentTag g :: *
    gameScore    :: g -> AgentTag g -> Int
    gameMoves    :: g -> [g]
    gameNullMove :: g -> g

instance Game Reversi where
    type AgentTag Reversi = Color

    gameScore rev me = relScoreFor me (reversiBoard rev)
    
    gameMoves rev =
        validMoves (reversiPlayer rev) (reversiBoard rev)
        |> fmap (Reversi $ other $ reversiPlayer rev)
    
    gameNullMove = reversiNullMove

type Agent g = AgentTag g -> g -> IO (Maybe g)

greedyAI :: Game g => Agent g
greedyAI me game = pure $
    gameMoves game
    |> maximumOn greedy
    where
        greedy g = gameScore g me - gameScore game me

type Alpha = Int
type Beta  = Int
type Depth = Int
type Score = Int

minimaxPruneAI :: forall g. Game g => Int -> Agent g
minimaxPruneAI maxDepth me game = pure $
    gameMoves game
    |> maximumOn (minPlay minBound maxBound maxDepth)
    where
        maxPlay, minPlay :: Alpha -> Beta -> Depth -> g -> Score

        maxPlay alpha beta depth game
            | depth <= 0            = gameScore game me
            | null (gameMoves game) = minPlay alpha beta (depth-1) (gameNullMove game)
            | otherwise             = go minBound alpha (gameMoves game)
                where
                    go value alpha [] = value
                    go value alpha (g:gs) =
                        if value' >= beta
                            then value'
                            else go value' alpha' gs
                        where
                            value' = max value (minPlay alpha beta (depth-1) g)
                            alpha' = max alpha value'
        
        minPlay alpha beta depth game
            | depth <= 0            = gameScore game me
            | null (gameMoves game) = maxPlay alpha beta (depth-1) (gameNullMove game)
            | otherwise             = go maxBound beta (gameMoves game)
                where
                    go value beta [] = value
                    go value beta (g:gs) =
                        if value' <= alpha
                            then value'
                            else go value' beta' gs
                        where
                            value' = min value (maxPlay alpha beta (depth-1) g)
                            beta'  = min beta value'

human :: Agent Reversi
human me game = do
    putStrLn ("You are: " ++ show me)
    ij <- readMove
    pure $ Just $ reversiMove ij game
    where
        readMove = do
            putStr ("Enter a move: ")
            [i, j] <- map read . words <$> getLine
            if (i, j) `elem` fmap fst (reversiValidPositions game)
                then pure (i, j)
                else do
                    putStrLn ("Invalid move.")
                    readMove

reversi :: Agent Reversi -> Agent Reversi -> IO ()
reversi p1 p2 = go False (cycle [p1, p2]) reversiInit
    where
        go skip (p:ps) game = do
            print game
            print (reversiScore game)
            putStrLn ""
            maybeGame' <- p (reversiPlayer game) game
            case maybeGame' of
                Just game' -> go False ps game'
                Nothing    ->
                    if skip then pure ()
                            else go True ps (reversiNullMove game)

reversiMinimaxFight d1 d2 = reversi (minimaxPruneAI d1) (minimaxPruneAI d2)

manual = reversi human human

main :: IO ()
main = do
    [d1, d2] <- map read <$> getArgs
    reversiMinimaxFight d1 d2