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: srces = self._get_sources() if len(srces) > 0: src = srces[self._src_idx] self._draw.draw_select(self._board, 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._board, 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