zeropod/plugin/chess/__init__.py

from os import path
from enum import Enum, auto
import time
import re
import itertools
from icecream import ic
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_SRC = auto()
    CHOOSE_DST = 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], player_color)

    state = GameState.MAIN_MENU
    menu_index = 0
    move = None
    score = None
    last_move = "Begin!"
    all_moves = dict()
    src_idx = 0
    dst_idx = 0
    src = ""

    while True:
        key = None
        if state == GameState.MAIN_MENU:
            draw.draw_menu(menu_index)
            key = cinput.get_one_shot()
        elif state == GameState.THINKING:
            draw.draw_thinking(last_move)
            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], player_color)
            last_move = move_str(move, player_color == "w")
            all_moves = dict()
            state = GameState.CHOOSE_SRC
        elif state == GameState.CHOOSE_SRC:
            if len(list(all_moves)) == 0:
                all_moves = get_all_moves(hist[-1])
            src = list(all_moves)[src_idx]
            draw.draw_select(last_move, list(all_moves)[src_idx])
            key = cinput.get_one_shot()
        elif state == GameState.CHOOSE_DST:
            src = list(all_moves)[src_idx]
            draw.draw_select(last_move, src, all_moves[src][dst_idx])
            key = cinput.get_one_shot()
        elif state == GameState.GAME_OVER:
            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_SRC
                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:
                        player_color = "b"
                        state = GameState.THINKING
                    else:
                        player_color = "w"
                        state = GameState.CHOOSE_SRC
                    hist = [pos]
                    draw.draw_board(hist[-1] if player_color == "w" else hist[-1].rotate(), player_color)
        elif state == GameState.CHOOSE_SRC:
            if key == Button.DIR_D or key == Button.DIR_R:
                src_idx += 1
                if src_idx >= len(all_moves.keys()):
                    src_idx = 0
            elif key == Button.DIR_U or key == Button.DIR_L:
                src_idx -= 1
                if src_idx < 0:
                    src_idx = len(list(all_moves)) - 1
            elif key == Button.BTN_A:
                state = GameState.MAIN_MENU
            elif key == Button.BTN_B:
                state = GameState.CHOOSE_DST
        elif state == GameState.CHOOSE_DST:
            if key == Button.DIR_D or key == Button.DIR_R:
                dst_idx += 1
                if dst_idx >= len(all_moves[src]):
                    dst_idx = 0
            elif key == Button.DIR_U or key == Button.DIR_L:
                dst_idx -= 1
                if dst_idx < 0:
                    dst_idx = len(all_moves[src]) - 1
            elif key == Button.BTN_A:
                state = GameState.CHOOSE_SRC
            elif key == Button.BTN_B:
                dst = all_moves[src][dst_idx]
                last_move = src + " - " + dst
                sfmove = sunfish.parse(src), sunfish.parse(dst)
                hist.append(hist[-1].move(sfmove))
                draw.draw_board(hist[-1].rotate(), player_color)
                state = GameState.THINKING
        else:
            if key == Button.BTN_A or key == Button.BTN_B:
                state = GameState.MAIN_MENU

def get_all_moves(pos: sunfish.Position):
    all_moves = dict()
    for src, dst in pos.gen_moves():
        all_moves.setdefault(sunfish.render(src), []).append(sunfish.render(dst))
    return all_moves

def move_str(move, shift = False):
    s = -119 if shift else 0
    return sunfish.render(s + move[0]) + " - " + sunfish.render(s + move[1])

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))