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

193 lines
7.3 KiB
Python

from os import path
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")
MAX_THINKING_SECONDS = 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):
self._board = Board() if fen == None else Board(fen)
self._last_move = "Begin!"
self._all_moves = dict()
self._src_idx = 0
self._dst_idx = 0
self._player_color = 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 run(self):
# either load the save game or start a new one
self._load_saved_or_init()
while True:
key = None
src = None
################
# HANDLE STATE #
################
# draw menu
if self._state == GameState.MAIN_MENU:
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_thinking(self._last_move)
# TODO: check game over
self._move = self._sf.get_best_move()
if self._move != None:
self._board.push(Move.from_uci(self._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
self._draw.draw_board(self._board, self._player_color)
# user picks source piece
elif self._state == GameState.CHOOSE_SRC:
if len(list(self._all_moves)) == 0:
for src, dst in map(lambda m: [Move.uci(m)[0:2], Move.uci(m)[2:4]], self._board.legal_moves):
self._all_moves.setdefault(src, []).append(dst)
if len(list(self._all_moves > 0)):
src = list(self._all_moves)[self._src_idx]
self._draw.draw_select(self._last_move, src)
key = self._cinput.get_one_shot()
else:
self._state = GameState.GAME_OVER
# user picks dest square
elif self._state == GameState.CHOOSE_DST:
src = list(self._all_moves)[self._src_idx]
dst = self._all_moves[src][self._dst_idx]
self._draw.draw_select(self._last_move, src, dst)
key = self._cinput.get_one_shot()
# game has ended
elif self._state == GameState.GAME_OVER:
key = self._cinput.get_one_shot()
################
# HANDLE INPUT #
################
# 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
self._init_new_game()
if self._menu_index == 2:
self._player_color = WHITE
self._state = GameState.CHOOSE_SRC
else:
self._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
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._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:
dst = self._all_moves[src][self._dst_idx]
self._move = src + dst
self._board.append(Move.from_uci(self._move))
self._draw.draw_board(self._board, self._player_color)
# TODO: check game over
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