228 lines
8.7 KiB
Python
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._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
|