content: {-# LANGUAGE LambdaCase, TupleSections #-}
  
  module RandomChess
  ( placeKings
  , placePawns
  , placeRemaining
  , emptyBoard
  , toFen
  , ChessBoard
  , Square (..)
  , BoardState (..)
  , getBoard
  )
  where
  
  import Control.Monad.State (State, get, gets, put)
  import Data.List (find, sortBy)
  import System.Random (Random, RandomGen, StdGen, random, randomR)
  
  type Pos = (Char, Int)
  
  type ChessBoard = [(Square, Pos)]
  
  data PieceRank  = King | Queen | Rook | Bishop | Knight | Pawn
    deriving (Enum, Bounded, Show, Eq, Ord)
  data PieceColor = Black | White
    deriving (Enum, Bounded, Show, Eq, Ord)
  
  data Square = ChessPiece PieceRank PieceColor | EmptySquare
    deriving (Eq, Ord)
  
  type PieceCount = [(Square, Int)]
  
  data BoardState = BoardState { board :: ChessBoard , generator :: StdGen }
  
  instance Show Square where
    show (ChessPiece King   Black) = "♚"
    show (ChessPiece Queen  Black) = "♛"
    show (ChessPiece Rook   Black) = "♜"
    show (ChessPiece Bishop Black) = "♝"
    show (ChessPiece Knight Black) = "♞"
    show (ChessPiece Pawn   Black) = "♟"
    show (ChessPiece King   White) = "♔"
    show (ChessPiece Queen  White) = "♕"
    show (ChessPiece Rook   White) = "♖"
    show (ChessPiece Bishop White) = "♗"
    show (ChessPiece Knight White) = "♘"
    show (ChessPiece Pawn   White) = "♙"
    show EmptySquare               = " "
  
  instance Random PieceRank where
    randomR (a, b) g = case randomR (fromEnum a, fromEnum b) g of
      (x, g'') -> (toEnum x, g'')
    random = randomR (minBound, maxBound)
  
  instance Random PieceColor where
    randomR (a, b) g = case randomR (fromEnum a, fromEnum b) g of
      (x, g'') -> (toEnum x, g'')
    random = randomR (minBound, maxBound)
  
  fullBoard :: PieceCount
  fullBoard =
    [ (ChessPiece King Black  , 1)
    , (ChessPiece Queen Black , 1)
    , (ChessPiece Rook Black  , 2)
    , (ChessPiece Bishop Black, 2)
    , (ChessPiece Knight Black, 2)
    , (ChessPiece Pawn Black  , 8)
    , (ChessPiece King White  , 1)
    , (ChessPiece Queen White , 1)
    , (ChessPiece Rook White  , 2)
    , (ChessPiece Bishop White, 2)
    , (ChessPiece Knight White, 2)
    , (ChessPiece Pawn White  , 8)
    , (EmptySquare            , 32)
    ]
  
  emptyBoard :: ChessBoard
  emptyBoard = fmap (EmptySquare,) . (,) <$> ['a'..'h'] <*> [1..8]
  
  replaceSquareByPos :: (Square, Pos) -> ChessBoard -> ChessBoard
  replaceSquareByPos e@(_, p) = fmap (\x -> if p == snd x then e else x)
  
  isPosOccupied :: Pos -> ChessBoard -> Bool
  isPosOccupied p = occupied . find (\x -> p == snd x)
   where
    occupied (Just (EmptySquare, _)) = False
    occupied _                       = True
  
  isAdjacent :: Pos -> Pos -> Bool
  isAdjacent (x1, y1) (x2, y2) =
    let upOrDown    = (pred y1 == y2 || succ y1 == y2)
        leftOrRight = (pred x1 == x2 || succ x1 == x2)
    in  (x2 == x1 && upOrDown)
          || (pred x1 == x2 && upOrDown)
          || (succ x1 == x2 && upOrDown)
          || (leftOrRight && y1 == y2)
  
  fen :: Square -> String
  fen (ChessPiece King   Black) = "k"
  fen (ChessPiece Queen  Black) = "q"
  fen (ChessPiece Rook   Black) = "r"
  fen (ChessPiece Bishop Black) = "b"
  fen (ChessPiece Knight Black) = "n"
  fen (ChessPiece Pawn   Black) = "p"
  fen (ChessPiece King   White) = "K"
  fen (ChessPiece Queen  White) = "Q"
  fen (ChessPiece Rook   White) = "R"
  fen (ChessPiece Bishop White) = "B"
  fen (ChessPiece Knight White) = "N"
  fen (ChessPiece Pawn   White) = "P"
  
  boardSort :: (Square, Pos) -> (Square, Pos) -> Ordering
  boardSort (_, (x1, y1)) (_, (x2, y2)) | y1 < y2  = GT
                                        | y1 > y2  = LT
                                        | y1 == y2 = compare x1 x2
  
  toFen :: ChessBoard -> String
  toFen [] = " w - - 0 1" <> []
  toFen b = scanRow (fst <$> take 8 b) 0
    where
      scanRow [] 0               = nextRow
      scanRow [] n               = show n <> nextRow
      scanRow (EmptySquare:xs) n = scanRow xs (succ n)
      scanRow (x:xs) 0           = nextPiece x xs
      scanRow (x:xs) n           = show n <> nextPiece x xs
      nextRow = "/" <> toFen (drop 8 b)
      nextPiece x xs = fen x <> scanRow xs 0
  
  -- State functions
  withStateGen :: (StdGen -> (a, StdGen)) -> State BoardState a
  withStateGen f = do
    currentState <- get
    let gen1 = generator currentState
    let (x, gen2) = f gen1
    put (currentState {generator = gen2})
    pure x
  
  randomPos :: State BoardState Pos
  randomPos = do
    boardState <- gets board
    chr <- withStateGen (randomR ('a', 'h'))
    num <- withStateGen (randomR (1, 8))
    let pos = (chr, num)
    if isPosOccupied pos boardState then
      randomPos
    else
      pure pos
  
  randomPiece :: State BoardState Square
  randomPiece = ChessPiece <$> withStateGen random <*> withStateGen random
  
  placeKings :: State BoardState ()
  placeKings = do
    currentState <- get
    p1 <- randomPos
    p2 <- randomPos
    if p1 `isAdjacent` p2 || p1 == p2
      then placeKings
      else do
        let updatedBoard = replaceSquareByPos (ChessPiece King White, p1) $
              replaceSquareByPos (ChessPiece King Black, p2) (board currentState)
        put currentState { board = updatedBoard }
  
  placePawns :: State BoardState ()
  placePawns = withStateGen (randomR (1, 16)) >>= go
    where
      go :: Int -> State BoardState ()
      go 0 = pure ()
      go n = do
        currentState <- get
        pos <- randomPos
        color <- withStateGen random
        let pawn = ChessPiece Pawn color
        let currentBoard = board currentState
        if promoted color == snd pos || isPosOccupied pos currentBoard ||
                             enpassant color == snd pos || firstPos color == snd pos
          then go n
          else do
            put currentState { board = replaceSquareByPos (pawn, pos) currentBoard }
            go $ pred n
      promoted White = 8
      promoted Black = 1
      enpassant White = 5
      enpassant Black = 4
      firstPos White = 1
      firstPos Black = 8
  
  placeRemaining :: State BoardState ()
  placeRemaining =
    withStateGen (randomR (5, sum $ fmap snd remaining)) >>= go remaining
    where
      remaining = filter (\case
                           (ChessPiece King _, _) -> False
                           (ChessPiece Pawn _, _) -> False
                           (EmptySquare, _) -> False
                           _ -> True)
                          fullBoard
  
      go :: PieceCount -> Int -> State BoardState ()
      go _ 0 = pure ()
      go remaining n = do
        currentState <- get
        let currentBoard = board currentState
        position <- randomPos
        piece <- randomPiece
        if not (isPermitted piece) || isPosOccupied position currentBoard
          then go remaining n
          else do
            let updatedBoard = replaceSquareByPos (piece, position) currentBoard
            put currentState { board = updatedBoard }
            go (consume piece remaining) (pred n)
        where
          isPermitted p =
            case find ((==p) . fst) remaining of
              Just (_, count) -> count > 0
              Nothing         -> False
          consume p'' = fmap (\(p, c) -> if p == p'' then (p, pred c) else (p, c))
  
  getBoard :: State BoardState ChessBoard
  getBoard = gets (sortBy boardSort . board)
  