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: array[2, Location],
                   bishops: array[2, Location], knights: array[2, Location],
                   pawns: array[8, 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
    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))


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 &= "- - - - - - - -"
    #const indeces = [7, 6, 5, 4, 3, 2, 1, 0]
    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)
    # Initialize all positions to a known default state.
    # This is useful in newChessBoardFromFEN
    result.pieces.black.king = (-1, -1)
    result.pieces.white.king = (-1, -1)
    result.pieces.black.queen = (-1, -1)
    result.pieces.white.queen = (-1, -1)
    result.pieces.black.rooks = [(-1, -1), (-1, -1)]
    result.pieces.white.rooks = [(-1, -1), (-1, -1)]
    result.pieces.black.bishops = [(-1, -1), (-1, -1)]
    result.pieces.white.bishops = [(-1, -1), (-1, -1)]
    result.pieces.black.knights = [(-1, -1), (-1, -1)]
    result.pieces.white.knights = [(-1, -1), (-1, -1)]
    result.pieces.black.pawns = [(-1, -1), (-1, -1), (-1, -1), (-1, -1), (-1, -1), (-1, -1), (-1, -1), (-1, -1)]
    result.pieces.white.pawns = [(-1, -1), (-1, -1), (-1, -1), (-1, -1), (-1, -1), (-1, -1), (-1, -1), (-1, -1)]
    # 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:
                                        # Find first empty slot in the pieces array
                                        for i, e in result.pieces.black.pawns:
                                            if e == (-1, -1):
                                                result.pieces.black.pawns[i] = (row, column)
                                                break
                                    of Bishop:
                                        if result.pieces.black.bishops[0] == (-1, -1):
                                            result.pieces.black.bishops[0] = (row, column)
                                        else:
                                            result.pieces.black.bishops[1] = (row, column)
                                    of Knight:
                                        if result.pieces.black.knights[0] == (-1, -1):
                                            result.pieces.black.knights[0] = (row, column)
                                        else:
                                            result.pieces.black.knights[1] = (row, column)
                                    of Rook:
                                        if result.pieces.black.rooks[0] == (-1, -1):
                                            result.pieces.black.rooks[0] = (row, column)
                                        else:
                                            result.pieces.black.rooks[1] = (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:
                                        for i, e in result.pieces.white.pawns:
                                            if e == (-1, -1):
                                                result.pieces.white.pawns[i] = (row, column)
                                                break
                                    of Bishop:
                                        if result.pieces.white.bishops[0] == (-1, -1):
                                            result.pieces.white.bishops[0] = (row, column)
                                        else:
                                            result.pieces.white.bishops[1] = (row, column)
                                    of Knight:
                                        if result.pieces.white.knights[0] == (-1, -1):
                                            result.pieces.white.knights[0] = (row, column)
                                        else:
                                            result.pieces.white.knights[1] = (row, column)
                                    of Rook:
                                        if result.pieces.white.rooks[0] == (-1, -1):
                                            result.pieces.white.rooks[0] = (row, column)
                                        else:
                                            result.pieces.white.rooks[1] = (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"

    echo "Running tests"
    var b = newDefaultChessboard()
    # Ensure pawns are in the correct location
    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"