chess/scripts/simulation.py

import uuid
from binding.python_c_ffi import Board
from binding.python_c_ffi import Move
from binding.python_c_ffi import ChessFFI
from binding.python_c_ffi import WHITE
from binding.python_c_ffi import print_board
from binding.python_c_ffi import sq_to_coord
from scripts.evaluation import RandomEval
from scripts.evaluation import NegaMaxEval
from scripts.format import LongPGNFormatter


YMD_HM = "%Y-%m-%d-%H-%M"
MAX_MOVES = 256
DATA_PATH = "./data/games"


class Engine:
    def __init__(self, strat_white=None, strat_black=None, max_plys=500):
        self.chess_ffi = ChessFFI()
        self.board = Board()
        self.max_plys = max_plys
        self.moves = []
        self.fens = []
        self._load_attack_cache()
        self._seed_engine()

        random_strat = RandomEval(chess_ffi=self.chess_ffi)

        if not strat_white:
            self.strat_white = random_strat
        else:
            self.strat_white = strat_white

        if not strat_black:
            self.strat_black = random_strat
        else:
            self.strat_black = strat_black

        self.strategies = {
            "white": {
                "name": self.strat_white.NAME,
                "params": self.strat_white.get_params()
            },
            "black": {
                "name": self.strat_black.NAME,
                "params": self.strat_black.get_params()
            },
        }


    def run(self, fen, save=True):
        self.chess_ffi.load_fen(self.board, fen)
        self._clear_moves()
        plys = 0

        self.fens.append(fen)
        while plys <= self.max_plys:
            print_board(self.board)

            moves_buf = (Move * MAX_MOVES)()
            n_legal = self.chess_ffi.get_legal_moves(self.board, moves_buf)

            game_over, result = self._is_game_over(n_legal)
            if game_over:
                print(result)
                break

            legal = [moves_buf[i] for i in range(n_legal)]
            if self.board.side_to_move == WHITE:
                best_move = self.strat_white.get_best_move(self.board, legal)
            else:
                best_move = self.strat_black.get_best_move(self.board, legal)

            new_board = Board()
            if not self.chess_ffi.apply_move_on_copy(self.board, new_board, best_move):
                print("ERROR: apply_move_on_copy failed")
                break

            self.board = new_board
            plys += 1

            fen_string = self.chess_ffi.board_to_fen(self.board)
            self.fens.append(fen_string)
            
            move = self.to_uci(best_move)
            self.moves.append(move)

            if self.board.halfmove_clock >= 100:
                print("draw (50-move rule)")
                result = "1/2-1/2"
                break

        print("done")
        print(len(self.moves))

        if save:
            print("saving game to disk.")
            id = uuid.uuid4().hex

            formatter = LongPGNFormatter(
                save_path=DATA_PATH,
                strategies=self.strategies,
                result=result
            )
            formatter.save_game(id, self.moves, self.fens)


    def side_tag(self, side):
        return 'w' if side == WHITE else 'b'


    def load_fen(self, fen):
        self.chess_ffi.load_fen(self.board, fen)


    def to_uci(self, move):
        fr = sq_to_coord(getattr(move, "from"))
        to = sq_to_coord(move.to)

        # Only add a promotion letter if promo is set
        promo_letter = ""
        promo_val = int(getattr(move, "promo", 0) or 0)
        if promo_val:
            # Normalize piece id to type 0..5 (P,N,B,R,Q,K)
            pt = promo_val % 6
            # Map N=1, B=2, R=3, Q=4
            letter_map = {1: "n", 2: "b", 3: "r", 4: "q"}
            promo_letter = letter_map.get(pt, "")

        return f"{fr}{to}{promo_letter}"


    def _load_attack_cache(self):
        self.chess_ffi.init_attack_cache()


    def _clear_moves(self):
        self.moves = []


    def _seed_engine(self):
        import ctypes as C, time

        seed = time.time_ns()
        s32 = int(seed) & 0xFFFFFFFF

        libc = C.CDLL("libc.so.6")   # on Ubuntu
        libc.srand.argtypes = (C.c_uint,)
        libc.srand.restype  = None
        libc.srand(C.c_uint(s32))
        return s32


    def _is_game_over(self, legal_moves):
        if legal_moves == 0:
            if self.chess_ffi.in_check(self.board, self.board.side_to_move):
                # Checkmate.
                if self.board.side_to_move == WHITE:
                    result = "0-1"
                else:
                    result = "1-0"
            else:
                # Stalemate.
                result = "1/2-1/2"
            return True, result
        return False, None
    

if __name__ == "__main__":
    chess_ffi = ChessFFI()
    nega_strat_1 = NegaMaxEval(chess_ffi=chess_ffi, depth=2, cp_window=5)
    nega_strat_2 = NegaMaxEval(chess_ffi=chess_ffi, depth=1, cp_window=5)


    engine = Engine(
        strat_white=nega_strat_1,
        strat_black=nega_strat_2,
    )

    fen = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
    engine.run(fen, save=True)