src/chess_tutor/classical_eval/board.py

"""Board representation and FEN parsing."""

from __future__ import annotations

from dataclasses import dataclass, field

from . import bitboards as bb
from .types import (
    BISHOP,
    BLACK,
    BLACK_BISHOP,
    BLACK_KING,
    BLACK_KNIGHT,
    BLACK_PAWN,
    BLACK_QUEEN,
    BLACK_ROOK,
    KING,
    KNIGHT,
    PAWN,
    QUEEN,
    ROOK,
    WHITE,
    WHITE_BISHOP,
    WHITE_KING,
    WHITE_KNIGHT,
    WHITE_PAWN,
    WHITE_QUEEN,
    WHITE_ROOK,
    make_piece,
)

PIECE_FROM_CHAR = {
    "P": (PAWN, WHITE),
    "N": (KNIGHT, WHITE),
    "B": (BISHOP, WHITE),
    "R": (ROOK, WHITE),
    "Q": (QUEEN, WHITE),
    "K": (KING, WHITE),
    "p": (PAWN, BLACK),
    "n": (KNIGHT, BLACK),
    "b": (BISHOP, BLACK),
    "r": (ROOK, BLACK),
    "q": (QUEEN, BLACK),
    "k": (KING, BLACK),
}


@dataclass
class Board:
    colours: list[int] = field(default_factory=lambda: [0, 0])
    pieces: list[int] = field(default_factory=lambda: [0] * 6)
    turn: int = WHITE
    psqtmat: int = 0


def board_from_fen(fen: str, psqt: list[list[int]] | None = None) -> Board:
    parts = fen.split()
    ranks = parts[0].split("/")
    board = Board()
    sq = 56
    for rank_str in ranks:
        file = 0
        for ch in rank_str:
            if ch.isdigit():
                file += int(ch)
            else:
                pt, colour = PIECE_FROM_CHAR[ch]
                s = sq + file
                board.colours[colour] |= 1 << s
                board.pieces[pt] |= 1 << s
                if psqt is not None:
                    board.psqtmat += psqt[make_piece(pt, colour)][s]
                file += 1
        sq -= 8
    board.turn = WHITE if (parts[1] if len(parts) > 1 else "w") == "w" else BLACK
    return board