zeropod/plugin/chess/__init__.py

from os import path
from enum import Enum, auto
import time
import re
import itertools
from cinput import ControlInput, Button
from graphics import Graphics
from .draw import Draw
from . import sunfish

class GameState(Enum):
    MAIN_MENU = auto()
    THINKING = auto()
    CHOOSE_ROW = auto()
    CHOOSE_COL = auto()
    GAME_OVER = auto()

SAVE_FILE = path.join("data", "chess_state.fen")
MAX_THINKING_SECONDS = 3

def execute(cinput: ControlInput, graphics: Graphics, _):
    graphics.clear()
    draw = Draw(graphics)
    hist: list[sunfish.Position]
    searcher = sunfish.Searcher()

    if path.exists(SAVE_FILE):
        with open(SAVE_FILE, "r") as fen:
            fenstr = fen.read()
            hist = [parseFEN(fenstr)]
    else:
        hist = [sunfish.Position(
                sunfish.initial, 0, (True,True), (True,True), 0, 0)]

    player_color = 'wb'[get_color(hist[-1])]
    draw.draw_board(hist[-1].board if player_color == 'w' else hist[-1].rotate().board)

    state = GameState.MAIN_MENU
    menu_index = 0
    move = None
    score = None
    last_move = "Begin!"

    while True:
        if state == GameState.MAIN_MENU:
            draw.draw_menu(menu_index)
        elif state == GameState.THINKING:
            draw.draw_thinking()
            start = time.time()
            for _, move, score in searcher.search(hist[-1], hist):
                if time.time() - start > MAX_THINKING_SECONDS:
                    break
            hist.append(hist[-1].move(move))
            draw.draw_board(hist[-1].board if player_color == 'b' else hist[-1].rotate().board)
            shift = -119 if player_color == 'b' else 0
            last_move = sunfish.render(shift+move[0]) + " -> " + sunfish.render(shift+move[1])
            if score == sunfish.MATE_UPPER:
                draw.draw_checkmate(last_move, False)
                state = GameState.GAME_OVER
            else:
                draw.draw_rowsel(last_move)
                state = GameState.CHOOSE_ROW

        key = cinput.get_one_shot()
        if state == GameState.MAIN_MENU:
            if key == Button.DIR_U:
                menu_index -= 1
                if menu_index < 0:
                    menu_index = 0
            elif key == Button.DIR_D:
                menu_index += 1
                if menu_index > 3:
                    menu_index = 3
            elif key == Button.BTN_B:
                if menu_index == 0:
                    state = GameState.CHOOSE_ROW
                elif menu_index == 1:
                    with open(SAVE_FILE, "w") as fen:
                        fen.write(renderFEN(hist[-1]))
                    return
                elif menu_index == 2 or menu_index == 3:
                    pos = sunfish.Position(
                            sunfish.initial, 0, (True,True), (True,True), 0, 0)
                    if menu_index == 3:
                        state = GameState.THINKING
                    else:
                        state = GameState.CHOOSE_ROW
                    hist = [pos]
                    draw.draw_board(pos.board)

        else:
            if key == Button.BTN_A:
                state = GameState.MAIN_MENU

def get_color(pos):
    return 1 if pos.board.startswith('\n') else 0

def parseFEN(fen):
    board, color, castling, enpas, _, _ = fen.split()
    board = re.sub(r'\d', (lambda m: '.'*int(m.group(0))), board)
    board = list(21*' ' + '  '.join(board.split('/')) + 21*' ')
    board[9::10] = ['\n']*12
    board = ''.join(board)
    wc = ('Q' in castling, 'K' in castling)
    bc = ('k' in castling, 'q' in castling)
    ep = sunfish.parse(enpas) if enpas != '-' else 0
    score = sum(sunfish.pst[p][i] for i,p in enumerate(board) if p.isupper())
    score -= sum(sunfish.pst[p.upper()][119-i] for i,p in enumerate(board) if p.islower())
    pos = sunfish.Position(board, score, wc, bc, ep, 0)
    return pos if color == 'w' else pos.rotate()

def renderFEN(pos, half_move_clock=0, full_move_clock=1):
    color = 'wb'[get_color(pos)]
    if color == 'b':
        pos = pos.rotate()
    board = '/'.join(pos.board.split())
    board = re.sub(r'\.+', (lambda m: str(len(m.group(0)))), board)
    castling = ''.join(itertools.compress('KQkq', pos.wc[::-1]+pos.bc)) or '-'
    ep = sunfish.render(pos.ep) if not pos.board[pos.ep].isspace() else '-'
    clock = '{} {}'.format(half_move_clock, full_move_clock)
    return ' '.join((board, color, castling, ep, clock))