module Main
where

import Data.Map (Map)
import qualified Data.Map as M
import Data.Bits
import Data.Char
import Control.Monad.State
import Data.List

data Color = Black | White deriving (Show, Eq)
type Board = Map Pos Piece
data CastleStatus = CS { 
      wk :: Bool, 
      wq :: Bool, 
      bk :: Bool,
      bq :: Bool
} deriving (Show)

type Chess = State GameState 

data GameState = GameState {
      board :: Board,
      turn :: Color,
      cursor :: Pos,
      castle :: CastleStatus
} deriving Show

data Piece = Piece {
      pieceType::PieceType, 
      pieceColor::Color, 
      location :: Pos
} deriving Eq

showPiece :: Piece -> String
showPiece p = convert letter ++ "@" ++ show (location p)
    where convert = if (pieceColor p == White) then (map toUpper) else (map toLower)
          letter = case (pieceType p) of 
                     Rook -> "r"
                     Knight -> "n"
                     Bishop -> "b"
                     Queen -> "q"
                     King -> "k"
                     Pawn -> "p"

instance Show Piece
    where show = showPiece

data PieceType = Rook | Knight | Bishop | King | Queen | Pawn deriving (Eq, Show)
type Pos = Int
type Move = (Pos, Pos)

genMoves :: Chess [Move]
genMoves = friendlyPieces >>= concatMapM genPieceMoves

concatMapM :: (Monad m) => (a -> m [b]) -> [a] -> m [b]
concatMapM f xs = liftM concat $ mapM f xs

genPieceMoves :: Piece -> Chess [Move]
genPieceMoves p = case pieceType p of
                    Pawn -> createPawnMoves (location p)
                    King -> liftM2 (++) createPieceMoves createCastleMoves
                    piece -> createPieceMoves
    where createPieceMoves = targets p >>= mapSource (location p)

mapSource :: Pos -> [Pos] -> Chess [Move]
mapSource = mapM . buildMove

targets :: Piece -> Chess [Pos]
targets piece = concatMapM ((moveFunction piece) (location piece)) (directions piece)

directions :: Piece -> [Pos]
directions (Piece Rook _ _) = straight
directions (Piece Bishop _ _) = diagonal
directions (Piece Knight _ _) = [33,31,-33,-31,18,14,-18,-14]
directions (Piece King _ _) = straight ++ diagonal
directions (Piece Queen _ _) = straight ++ diagonal
directions _ = []

straight, diagonal :: [Pos]
straight = [1, -1, 16, -16]
diagonal = [15, 17, negate 15, negate 17]

moveFunction :: Piece -> (Pos -> Pos -> Chess [Pos])
moveFunction (Piece Knight _ _) = step
moveFunction (Piece King _ _) = step
moveFunction _ = slide

if' :: Bool -> a -> a -> a
if' p t f = if p then t else f

step :: Pos -> Pos -> Chess [Pos]
step source target = do isEnemy <- unfriendly dest
                        return [dest | valid dest, isEnemy]
    where dest = source+target

friendly :: Pos -> Chess Bool
friendly pos = gets $ \s -> M.member pos (board s) && pieceColor ((board s) M.! pos) == (turn s)

unfriendly :: Pos -> Chess Bool
unfriendly = liftM not . friendly

slide :: Pos -> Int -> Chess [Pos]
slide source dir = takeWhileM' isEmpty unfriendly (ray source dir)

takeWhileM' :: Monad m => (a -> m Bool) -> (a -> m Bool) -> [a] -> m [a]
takeWhileM' _ _ [] = return []
takeWhileM' p q (x:xs) = caseM [
                          p x   --> liftM (x:) (takeWhileM' p q xs),
                          q x   --> return [x],
                          elseM --> return []
                         ]
(-->) = (,)
elseM :: (Monad m) => m Bool
elseM = return True
caseM :: (Monad m) => [(m Bool, m a)] -> m a
caseM ((p,v):pvs) = p >>= ($ caseM pvs) . flip if' v

buildMove :: Pos -> Pos -> Chess Move
buildMove x y = return (x, y)

friendlyPieces :: Chess [Piece]
friendlyPieces = gets $ \s -> M.elems $ M.filter ((turn s ==) . pieceColor) (board s)

createCastleMoves :: Chess [Move]
createCastleMoves = do isOk <- ksRequirements
                       if (isOk) then ksCastleMove else return []

{-
gets $ liftM2 (++) ksCastle qsCastle
    where ksCastle s = if (ksRequirements s) then (ksCastleMove s) else []
          qsCastle s = if (qsRequirements s) then (qsCastleMove s) else []
-}

qsRequirements :: Chess Bool
qsRequirements = gets $ \gs -> if (turn gs == White) 
                               then ((and $ (wq (castle gs)) : map (not . flip M.member (board gs)) [1,2,3])) 
                               else ((and $ (bq (castle gs)) : map (not . flip M.member (board gs)) [113,114,115]))

ksRequirements :: Chess Bool
ksRequirements = gets $ \gs -> if (turn gs == White)
                               then ((and $ (wk (castle gs)) : map (not . flip M.member (board gs)) [5,6]))
                               else ((and $ (bk (castle gs)) : map (not . flip M.member (board gs)) [117,118]))

ksCastleMove :: Chess [Move]
ksCastleMove = gets $ \gs -> if (turn gs == White)
                             then (sequence [buildMove 4 6])
                             else (sequence [buildMove 116 118])

qsCastleMove :: Chess [Move]
qsCastleMove = undefined --gs | turn gs == White = [buildMove 4 1]
--qsCastleMove = undefined --gs | turn gs == Black = [buildMove 116 112]

createPawnMoves :: Pos -> Chess [Move]
createPawnMoves pos = liftM concat . sequence . map ($ pos) $ [
                       leftPawnCaptures, 
                       rightPawnCaptures, 
                       singlePawnMove, 
                       doublePawnMove ]

leftPawnCaptures :: Pos -> Chess [Move]
leftPawnCaptures pos = gets $ \s -> if' (hasLeftCapture s pos) [(pos, pos+15)] []
    where hasLeftCapture s pos = and [
                                  M.member (pos+15) (board s),
                                  pieceColor ((board s) M.! (pos+15)) == turn s
                                 ]


rightPawnCaptures :: Pos -> Chess [Move]
rightPawnCaptures pos = gets $ \s -> if' (hasRightCapture s pos) [(pos, pos+17)] []
    where hasRightCapture s pos = and [
                                   M.member (pos+17) (board s),
                                   pieceColor ((board s) M.! (pos+17)) == turn s
                                  ]
                                   
singlePawnMove :: Pos -> Chess [Move]
singlePawnMove pos = gets $ \s -> if' (hasSingle s pos) [(pos, pos+16)] []
    where hasSingle s pos = and [
                             not (M.member (pos+16) (board s))
                            ]

doublePawnMove :: Pos -> Chess [Move]
doublePawnMove pos = gets $ \s -> if' (hasDouble s pos) [(pos, pos+32)] []
    where hasDouble s pos = and [
                             not (M.member (pos+32) (board s))
                            ]

makeMove :: Move -> Chess ()
makeMove (source, target) = modify $ \s -> s{
                              board = (remove source . copyFrom source target . remove target) (board s),
                              turn = oppositeColor (turn s)
                            }

remove :: Pos -> Board -> Board
remove pos board = if' (M.member pos board) (M.delete pos board) board

copyFrom :: Pos -> Pos -> Board -> Board
copyFrom source target board = M.insert target (board M.! source) board

oppositeColor :: Color -> Color
oppositeColor White = Black
oppositeColor Black = White

isEmpty :: Pos -> Chess Bool
isEmpty pos = gets $ not . M.member pos . board

valid, invalid :: Pos -> Bool
invalid = not.valid
valid = (== 0) . (0x88 .&.)

emptyBoard :: Board
emptyBoard = M.empty

emptyCS :: CastleStatus
emptyCS = CS False False False False

emptyGameState :: GameState
emptyGameState = GameState emptyBoard White 112 emptyCS

loadFEN :: String -> GameState
loadFEN = foldr ($) emptyGameState . zipWith ($) funcs . words
    where funcs = [parseBoard,parseTurn,parseCS]

parseBoard :: String -> GameState -> GameState
parseBoard = flip (foldl' place)

parseTurn :: String -> GameState -> GameState
parseTurn ('w':_) gs = gs{turn=White}
parseTurn ('b':_) gs = gs{turn=Black}

parseCS :: String -> GameState -> GameState
parseCS "" gs = gs
parseCS ('K':cs) gs = parseCS cs gs{ castle = (castle gs){wk=True}}
parseCS ('Q':cs) gs = parseCS cs gs{ castle = (castle gs){wq=True}}
parseCS ('k':cs) gs = parseCS cs gs{ castle = (castle gs){bk=True}}
parseCS ('q':cs) gs = parseCS cs gs{ castle = (castle gs){bq=True}}
parseCS ('-':cs) gs = gs


place :: GameState -> Char -> GameState
place state '/' = state{cursor = cursor state - 24}
place state c | isDigit c = state{cursor = cursor state + digitToInt c}
place state c | otherwise = state{board = M.insert (cursor state) piece (board state), 
                                  cursor=cursor state + 1}
              where piece = makePiece c (cursor state)

makePiece :: Char -> Pos -> Piece
makePiece c pos = Piece typ color pos
    where typ = case (toUpper c) of
                  'P' -> Pawn
                  'N' -> Knight
                  'B' -> Bishop
                  'R' -> Rook
                  'Q' -> Queen
                  'K' -> King
          color = if' (isUpper c) White Black

file :: Pos -> Int
file = (+1) . (`mod` 16)

rank :: Pos -> Int
rank = (+1) . (`div` 16)

ray :: Pos -> Int -> [Pos]
ray pos dir = takeWhile valid (iterate (+dir) (pos+dir))

testboard :: GameState
testboard = loadFEN "4k3/8/8/8/8/8/8/4K2R w K - 0 1"

separateBy :: Eq a => a -> [a] -> [[a]]
separateBy ch string = case dropWhile (== ch) string of
                         [] -> []
                         s:ss -> (s:word) : separateBy ch rest
                             where (word,rest) = break (== ch) ss

epdtest :: IO ()
epdtest = interact $ unlines . map (testPosition . getPosition) . lines

testPosition :: (String, [Int]) -> String
testPosition (fen, answers) = case (length . evalState genMoves . loadFEN) fen of
                                answer | answer == head answers -> fen ++ ": Correct!"
                                answer | otherwise -> fen ++ ": Wrong! Answer: " ++ show answer ++ ", correct: " ++ (show $ head answers)

getPosition :: String -> (String, [Int])
getPosition fen = (position, nums)
    where fields = (map init . separateBy ';') fen
          position = head fields
          nums = cleanup (drop 1 fields)
          cleanup = map (read . drop 3)

main = interact id