This repository has been archived on 2022-12-29. You can view files and clone it, but cannot push or open issues or pull requests.
zeropod/plugin/chess/game.py

228 lines
8.7 KiB
Python

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._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