CPG/src/Chess/board.nim

# Copyright 2023 Mattia Giambirtone & All Contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ../util/matrix

export matrix

import std/strutils
import std/strformat


type
    Location = tuple[row, col: int]
    Pieces = tuple[king: Location, queen: Location, rooks: seq[Location],
                   bishops: seq[Location], knights: seq[Location],
                   pawns: seq[Location]]
    PieceColor* = enum
        None = 0,
        White, 
        Black
    PieceKind* = enum
        Empty = '\0',    # No piece
        Bishop = 'b', 
        King = 'k'
        Knight = 'n',
        Pawn = 'p', 
        Queen = 'q',
        Rook = 'r', 
    Piece* = object
        color*: PieceColor
        kind*: PieceKind
        captured*: bool
    Position* = object
        piece*: Piece
        location*: Location
    ChessBoard* = ref object
        ## A chess board object
        grid: Matrix[Piece]
        # Currently active color
        turn: PieceColor
        # Number of half moves since
        # last piece capture or pawn movement.
        # Used for the 50-move rule
        halfMoveClock: int
        # Full move counter. Increments
        # every 2 ply
        fullMoveCount: int
        # Stores metadata for castling.
        castling: tuple[white, black: tuple[queen, king: bool]]
        # En passant target square (see https://en.wikipedia.org/wiki/En_passant)
        # If en passant is not possible, both the row and
        # column of the position will be set to -1
        enPassantSquare: Position
        # Locations of all pieces
        pieces: tuple[white: Pieces, black: Pieces]
        # Locations of all attacked squares
        attacked: tuple[white: seq[Location], black: seq[Location]]


# Initialized only once, copied every time
var empty: seq[Piece] = @[]
for _ in countup(0, 63):
    empty.add(Piece(kind: Empty, color: None))


proc countPieces*(self: ChessBoard, kind: PieceKind, color: PieceColor): int =
    ## Counts the number of pieces with
    ## the given color and type
    case color:
        of White:
            case kind:
                of Pawn:
                    return self.pieces.white.pawns.len()
                of Bishop:
                    return self.pieces.white.bishops.len()
                of Knight:
                    return self.pieces.white.knights.len()
                of Rook:
                    return self.pieces.white.rooks.len()
                of Queen:
                    return int(self.pieces.white.queen != (-1, -1))
                of King:
                    # There shall be only one, forever
                    return 1
                else:
                    discard
        of Black:
            case kind:
                of Pawn:
                    return self.pieces.black.pawns.len()
                of Bishop:
                    return self.pieces.black.bishops.len()
                of Knight:
                    return self.pieces.black.knights.len()
                of Rook:
                    return self.pieces.black.rooks.len()
                of Queen:
                    return int(self.pieces.black.queen != (-1, -1))
                of King:
                    return 1
                else:
                    discard        
        of None:
            raise newException(ValueError, "invalid piece type")


proc countPieces*(self: ChessBoard, piece: Piece): int =
    ## Returns the number of pieces on the board that
    ## are of the same type and color of the given piece
    return self.countPieces(piece.kind, piece.color)


func rankToColumn(rank: int): int =
    ## Converts a chess rank (1-indexed) 
    ## into a 0-indexed column value for our 
    ## board. This converter is necessary because
    ## chess positions are indexed differently with
    ## respect to our internal representation
    const indeces = [7, 6, 5, 4, 3, 2, 1, 0]
    return indeces[rank - 1]


proc algebraicToPosition*(s: string): Location {.inline.} =
    ## Converts a square location from algebraic
    ## notation to its corresponding row and column
    ## in the chess grid (0 indexed)
    if len(s) != 2:
        raise newException(ValueError, "algebraic position must be of length 2")

    var s = s.toLowerAscii()
    if s[0] notin 'a'..'h':
        raise newException(ValueError, &"algebraic position has invalid first character ('{s[0]}')")
    if s[1] notin '1'..'8':
        raise newException(ValueError, &"algebraic position has invalid second character ('{s[1]}')")

    let file = int(uint8(s[0]) - uint8('a'))
    # Convert the rank character to a number
    let rank = rankToColumn(int(uint8(s[1]) - uint8('0')))
    return (rank, file)


proc getPiece*(self: ChessBoard, square: string): Piece =
    ## Gets the piece on the given square
    ## in algebraic notation
    let loc = square.algebraicToPosition()
    return self.grid[loc.row, loc.col]


proc `$`*(self: ChessBoard): string =
    result &= "- - - - - - - -"
    for i, row in self.grid:
        result &= "\n"
        for piece in row:
            if piece.kind == Empty:
                result &= "  "
                continue
            if piece.color == White:
                result &= &"{char(piece.kind).toUpperAscii()} "
            else:
                result &= &"{char(piece.kind)} "
        result &= &"{rankToColumn(i + 1) + 1}"
    result &= "\n- - - - - - - -"
    result &= "\na b c d e f g h"


proc newChessboard: ChessBoard =
    ## Returns a new, empty chessboard
    new(result)
    # Turns our flat sequence into an 8x8 grid
    result.grid = newMatrixFromSeq[Piece](empty, (8, 8))
    result.attacked = (@[], @[])
    result.enPassantSquare = Position(piece: Piece(kind: Empty, color: None), location: (-1, -1))


proc newChessboardFromFEN*(state: string): ChessBoard =
    ## Initializes a chessboard with the
    ## state encoded by the given FEN string
    result = newChessboard()
    var
        # Current location in the grid
        row = 0
        column = 0
        # Current section in the FEN string
        section = 0
        # Current index into the FEN string
        index = 0
        # Temporary variable to store the piece
        piece: Piece
    # See https://en.wikipedia.org/wiki/Forsyth%E2%80%93Edwards_Notation
    while index <= state.high():
        var c = state[index]
        if c == ' ':
            # Next section
            inc(section)
            inc(index)
            continue
        case section:
            of 0:
                # Piece placement data
                case c.toLowerAscii():
                    # Piece
                    of 'r', 'n', 'b', 'q', 'k', 'p':
                        # We know for a fact these values are in our 
                        # enumeration, so all is good
                        {.push.}
                        {.warning[HoleEnumConv]:off.}
                        piece = Piece(kind: PieceKind(c.toLowerAscii()), color: if c.isUpperAscii(): White else: Black)
                        {.pop.}
                        case piece.color:
                            of Black:
                                case piece.kind:
                                    of Pawn:
                                        result.pieces.black.pawns.add((row, column))
                                    of Bishop:
                                        result.pieces.black.bishops.add((row, column))
                                    of Knight:
                                        result.pieces.black.knights.add((row, column))
                                    of Rook:
                                        result.pieces.black.rooks.add((row, column))
                                    of Queen:
                                        result.pieces.black.queen = (row, column)
                                    of King:
                                        result.pieces.black.king = (row, column)
                                    else:
                                        discard
                            of White:
                                case piece.kind:
                                    of Pawn:
                                        result.pieces.white.pawns.add((row, column))
                                    of Bishop:
                                        result.pieces.white.bishops.add((row, column))
                                    of Knight:
                                        result.pieces.white.knights.add((row, column))
                                    of Rook:
                                        result.pieces.white.rooks.add((row, column))
                                    of Queen:
                                        result.pieces.white.queen = (row, column)
                                    of King:
                                        result.pieces.white.king = (row, column)
                                    else:
                                        discard
                            else:
                                discard
                        result.grid[row, column] = piece
                        inc(column)
                    of '/':
                        # Next row
                        inc(row)
                        column = 0
                    of '0'..'9':
                        # Skip x columns
                        let x = int(uint8(c) - uint8('0')) - 1
                        if x > 7:
                            raise newException(ValueError, "invalid skip value (> 8) in FEN string")
                        column += x
                    else:
                        raise newException(ValueError, "invalid piece identifier in FEN string")
            of 1:
                # Active color
                case c:
                    of 'w':
                        result.turn = White
                    of 'b':
                        result.turn = Black
                    else:
                        raise newException(ValueError, "invalid active color identifier in FEN string")
            of 2:
                # Castling availability
                case c:
                    of '-':
                        # Neither side can castle anywhere: do nothing,
                        # as the castling metadata is set to this state
                        # by default
                        discard
                    of 'K':
                        result.castling.white.king = true
                    of 'Q':
                        result.castling.white.queen = true
                    of 'k':
                        result.castling.black.king = true
                    of 'q':
                        result.castling.black.queen = true
                    else:
                        raise newException(ValueError, "invalid castling availability in FEN string")
            of 3:
                # En passant target square
                case c:
                    of '-':
                        # Field is already uninitialized to the correct state
                        discard
                    else:
                        result.enPassantSquare.location = state[index..index+1].algebraicToPosition()
                        # Just for cleanliness purposes, we fill in the other positional metadata as
                        # well
                        result.enPassantSquare.piece.color = if result.turn == Black: White else: Black
                        result.enPassantSquare.piece.kind = Pawn
                        # Square metadata is 2 bytes long
                        inc(index)
            of 4:
                # Halfmove clock
                var s = ""
                while not state[index].isSpaceAscii():
                    s.add(state[index])
                    inc(index)
                # Backtrack so the space is seen by the
                # next iteration of the loop
                dec(index)
                result.halfMoveClock = parseInt(s)
            of 5:
                # Fullmove number
                var s = ""
                while index <= state.high():
                    s.add(state[index])
                    inc(index)
                result.fullMoveCount = parseInt(s)
            else:
                raise newException(ValueError, "too many fields in FEN string")
        inc(index)


proc newDefaultChessboard*: ChessBoard =
    ## Initializes a chessboard with the
    ## starting position
    return newChessboardFromFEN("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1")


when isMainModule:
    proc testPiece(piece: Piece, kind: PieceKind, color: PieceColor) =
        doAssert piece.kind == kind and piece.color == color, &"expected piece of kind {kind} and color {color}, got {piece.kind} / {piece.color} instead"

    proc testPieceCount(board: ChessBoard, kind: PieceKind, color: PieceColor, count: int) =
        let pieces = board.countPieces(kind, color)
        doAssert pieces == count, &"expected {count} pieces of kind {kind} and color {color}, got {pieces} instead"

    echo "Running tests"
    var b = newDefaultChessboard()
    # Ensure correct number of pieces
    testPieceCount(b, Pawn, White, 8)
    testPieceCount(b, Pawn, Black, 8)
    testPieceCount(b, Knight, White, 2)
    testPieceCount(b, Knight, Black, 2)
    testPieceCount(b, Bishop, White, 2)
    testPieceCount(b, Bishop, Black, 2)
    testPieceCount(b, Rook, White, 2)
    testPieceCount(b, Rook, Black, 2)
    testPieceCount(b, Queen, White, 1)
    testPieceCount(b, Queen, Black, 1)
    testPieceCount(b, King, White, 1)
    testPieceCount(b, King, Black, 1)
    

    # Ensure pieces are in the correct location

    # Pawns
    for loc in ["a2", "b2", "c2", "d2", "e2", "f2", "g2", "h2"]:
        testPiece(b.getPiece(loc), Pawn, White)
    for loc in ["a7", "b7", "c7", "d7", "e7", "f7", "g7", "h7"]:
        testPiece(b.getPiece(loc), Pawn, Black)
    # Rooks
    testPiece(b.getPiece("a1"), Rook, White)
    testPiece(b.getPiece("h1"), Rook, White)
    testPiece(b.getPiece("a8"), Rook, Black)
    testPiece(b.getPiece("h8"), Rook, Black)
    # Knights
    testPiece(b.getPiece("b1"), Knight, White)
    testPiece(b.getPiece("g1"), Knight, White)
    testPiece(b.getPiece("b8"), Knight, Black)
    testPiece(b.getPiece("g8"), Knight, Black)
    # Bishops
    testPiece(b.getPiece("c1"), Bishop, White)
    testPiece(b.getPiece("f1"), Bishop, White)
    testPiece(b.getPiece("c8"), Bishop, Black)
    testPiece(b.getPiece("f8"), Bishop, Black)
    # Kings
    testPiece(b.getPiece("e1"), King, White)
    testPiece(b.getPiece("e8"), King, Black)
    # Queens
    testPiece(b.getPiece("d1"), Queen, White)
    testPiece(b.getPiece("d8"), Queen, Black)
    echo "All tests were successful"