zeropod/plugin/chess/game.py

from os import path
from typing import Optional
from enum import Enum, auto
import shutil
from icecream import ic
from cinput import ControlInput, Button
from graphics import Graphics
from .draw import Draw
from chess import Board, Move, WHITE, BLACK
from stockfish import Stockfish

class GameState(Enum):
    MAIN_MENU = auto()
    THINKING = auto()
    CHOOSE_SRC = auto()
    CHOOSE_DST = auto()
    GAME_OVER = auto()

class ChessGame:
    SAVE_FILE = path.join("data", "chess_state.fen")
    STOCKFISH_MOVE_TIME = 3

    def __init__(self, cinput: ControlInput, graphics: Graphics):
        self._sf = Stockfish(shutil.which("stockfish") or "stockfish")
        self._cinput = cinput
        self._draw = Draw(graphics)
        self._state = GameState.MAIN_MENU
        self._menu_index = 0

    def _init_new_game(self, fen = None, player_color = None):
        self._board = Board() if fen == None else Board(fen)
        self._move = "Begin!"
        self._all_moves = dict()
        self._src_idx = 0
        self._dst_idx = 0
        self._player_color = player_color if player_color != None else self._board.turn
        self._sf.set_fen_position(self._board.fen())
        self._draw.draw_board(self._board, self._player_color)

    def _load_saved_or_init(self):
        if path.exists(self.SAVE_FILE):
            with open(self.SAVE_FILE, "r") as fen:
                fen = fen.read()
                self._init_new_game(fen)
        else:
            self._init_new_game()

    def _save_game(self):
        with open(self.SAVE_FILE, "w") as fen:
            fen.write(self._board.fen())

    def _make_move(self, move: str):
        self._move = move
        self._board.push(Move.from_uci(move))
        self._sf.make_moves_from_current_position([move])
        self._draw.draw_board(self._board, self._player_color)

    def _square_sort_key(self, square: str):
        # convert to number between 00 to 88
        key = (int(square[1]) - 1) * 10 + ord(square[0]) - 97
        return key if self._player_color == WHITE else 88 - key

    def _get_sources(self):
        if len(list(self._all_moves)) == 0:
            moves = self._board.generate_legal_moves()
            mvarray = map(lambda m: [Move.uci(m)[0:2], Move.uci(m)[2:4]], moves)
            for src, dst in list(mvarray):
                self._all_moves.setdefault(src, []).append(dst)
        if len(list(self._all_moves)) > 0:
            return sorted(
                    list(self._all_moves),
                    key=lambda x: self._square_sort_key(x))
        return list[str]()

    def _get_dests(self):
        if self._src_idx < len(list(self._all_moves)):
            src = self._get_sources()[self._src_idx]
            return sorted(
                    self._all_moves[src],
                    key=lambda x: self._square_sort_key(x))
        return list[str]()

    def run(self):
        # either load the save game or start a new one
        self._load_saved_or_init()

        while True:
            key: Optional[Button] = None
            src: Optional[str] = None

            ################
            # HANDLE STATE #
            ################

            if self._board.is_game_over() and self._state != GameState.MAIN_MENU:
                self._state = GameState.GAME_OVER

            # draw menu
            if self._state == GameState.MAIN_MENU:
                self._draw.draw_board(self._board, self._player_color)
                self._draw.draw_menu(self._menu_index)
                key = self._cinput.get_one_shot()

            # computer makes a move
            elif self._state == GameState.THINKING:
                self._draw.draw_board(self._board, self._player_color)
                self._draw.draw_thinking(self._move)
                move = self._sf.get_best_move_time(self.STOCKFISH_MOVE_TIME * 1000)
                if move != None:
                    self._make_move(move)
                    # reset user moves and move state to user's turn
                    self._all_moves = dict()
                    self._src_idx = 0
                    self._dst_idx = 0
                    self._state = GameState.CHOOSE_SRC
                else:
                    self._state = GameState.GAME_OVER

            # user picks source piece
            elif self._state == GameState.CHOOSE_SRC:
                if len(self._get_sources()) > 0:
                    src = self._get_sources()[self._src_idx]
                    self._draw.draw_select(self._move, self._player_color, src)
                    key = self._cinput.get_one_shot(0.1)
                else:
                    self._state = GameState.GAME_OVER

            # user picks dest square
            elif self._state == GameState.CHOOSE_DST:
                src = self._get_sources()[self._src_idx]
                dst = self._get_dests()[self._dst_idx]
                self._draw.draw_select(self._move, self._player_color, src, dst)
                key = self._cinput.get_one_shot(0.1)

            # game has ended
            else:
                self._draw.draw_board(self._board, self._player_color)
                self._draw.draw_game_over(self._board.outcome())
                key = self._cinput.get_one_shot()

            ################
            # HANDLE INPUT #
            ################

            if key == None:
                continue

            # handle user input on main menu
            if self._state == GameState.MAIN_MENU:

                # menu cursor up and down
                if key == Button.DIR_U:
                    self._menu_index -= 1
                    if self._menu_index < 0:
                        self._menu_index = 0
                elif key == Button.DIR_D:
                    self._menu_index += 1
                    if self._menu_index > 3:
                        self._menu_index = 3

                # select current menu item
                elif key == Button.BTN_B:
                    if self._menu_index == 0: # play current game
                        self._state = GameState.CHOOSE_SRC
                    elif self._menu_index == 1: # save and quit
                        self._save_game()
                        return
                    elif self._menu_index == 2 or self._menu_index == 3: # new game
                        if self._menu_index == 2:
                            self._init_new_game(player_color=WHITE)
                            self._state = GameState.CHOOSE_SRC
                        else:
                            self._init_new_game(player_color=BLACK)
                            self._state = GameState.THINKING

            # handle user input when selecting source piece
            # TODO: some kind of cursor indicator on board
            elif self._state == GameState.CHOOSE_SRC:
                # move between source pieces
                cur_src = self._get_sources()[self._src_idx]

                if key == Button.DIR_D or key == Button.DIR_R:
                    self._src_idx += 1
                    if self._src_idx >= len(self._all_moves.keys()):
                        self._src_idx = 0
                elif key == Button.DIR_U or key == Button.DIR_L:
                    self._src_idx -= 1
                    if self._src_idx < 0:
                        self._src_idx = len(list(self._all_moves)) - 1

                # back out to main menu
                elif key == Button.BTN_A:
                    self._state = GameState.MAIN_MENU

                # move to destination picker
                elif key == Button.BTN_B:
                    self._dst_idx = 0
                    self._state = GameState.CHOOSE_DST

            # handle user input when selecting dest piece
            elif self._state == GameState.CHOOSE_DST:
                # move between dest squares for the given source
                if key == Button.DIR_D or key == Button.DIR_R:
                    self._dst_idx += 1
                    if self._dst_idx >= len(self._all_moves[src]):
                        self._dst_idx = 0
                elif key == Button.DIR_U or key == Button.DIR_L:
                    self._dst_idx -= 1
                    if self._dst_idx < 0:
                        self._dst_idx = len(self._all_moves[src]) - 1

                # back out to choose a different source piece
                elif key == Button.BTN_A:
                    self._dst_idx = 0
                    self._state = GameState.CHOOSE_SRC

                # make the move
                elif key == Button.BTN_B:
                    src = self._get_sources()[self._src_idx]
                    dst = self._get_dests()[self._dst_idx]
                    self._make_move(src + dst)
                    self._state = GameState.THINKING

            # handle user input on game over
            else:
                if key == Button.BTN_A or key == Button.BTN_B:
                    self._state = GameState.MAIN_MENU