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
Raw Normal View History

2022-04-25 14:49:42 -05:00
from os import path
from typing import Optional
2022-04-25 14:49:42 -05:00
from enum import Enum, auto
import shutil
from icecream import ic
2022-04-25 14:49:42 -05:00
from cinput import ControlInput, Button
from graphics import Graphics
from .draw import Draw
from chess import Board, Move, WHITE, BLACK
from stockfish import Stockfish
2022-04-25 14:49:42 -05:00
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")
2022-04-26 10:09:44 -05:00
STOCKFISH_MOVE_TIME = 3
2022-04-25 14:49:42 -05:00
def __init__(self, cinput: ControlInput, graphics: Graphics):
self._sf = Stockfish(shutil.which("stockfish") or "stockfish")
2022-04-25 14:49:42 -05:00
self._cinput = cinput
self._draw = Draw(graphics)
self._state = GameState.MAIN_MENU
self._menu_index = 0
2022-04-26 10:09:44 -05:00
def _init_new_game(self, fen = None, player_color = None):
self._board = Board() if fen == None else Board(fen)
2022-04-26 10:09:44 -05:00
self._move = "Begin!"
2022-04-25 14:49:42 -05:00
self._all_moves = dict()
self._src_idx = 0
self._dst_idx = 0
2022-04-26 10:09:44 -05:00
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)
2022-04-25 14:49:42 -05:00
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)
2022-04-25 14:49:42 -05:00
else:
self._init_new_game()
2022-04-25 14:49:42 -05:00
def _save_game(self):
with open(self.SAVE_FILE, "w") as fen:
fen.write(self._board.fen())
2022-04-25 14:49:42 -05:00
2022-04-26 10:09:44 -05:00
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
2022-04-30 11:00:14 -05:00
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))
2022-04-30 11:00:14 -05:00
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))
2022-04-30 11:00:14 -05:00
return list[str]()
2022-04-25 14:49:42 -05:00
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
2022-04-25 14:49:42 -05:00
################
# HANDLE STATE #
################
2022-04-30 11:00:14 -05:00
if self._board.is_game_over() and self._state != GameState.MAIN_MENU:
self._state = GameState.GAME_OVER
2022-04-25 14:49:42 -05:00
# draw menu
if self._state == GameState.MAIN_MENU:
self._draw.draw_board(self._board, self._player_color)
2022-04-25 14:49:42 -05:00
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)
2022-04-26 10:09:44 -05:00
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
2022-04-25 14:49:42 -05:00
# user picks source piece
elif self._state == GameState.CHOOSE_SRC:
2022-04-30 17:26:14 -05:00
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
2022-04-25 14:49:42 -05:00
# user picks dest square
elif self._state == GameState.CHOOSE_DST:
2022-04-30 11:00:14 -05:00
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)
2022-04-25 14:49:42 -05:00
# game has ended
2022-04-30 11:00:14 -05:00
else:
self._draw.draw_board(self._board, self._player_color)
2022-04-30 11:00:14 -05:00
self._draw.draw_game_over(self._board.outcome())
2022-04-25 14:49:42 -05:00
key = self._cinput.get_one_shot()
################
# HANDLE INPUT #
################
if key == None:
continue
2022-04-25 14:49:42 -05:00
# 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:
2022-04-26 10:09:44 -05:00
self._init_new_game(player_color=WHITE)
2022-04-25 14:49:42 -05:00
self._state = GameState.CHOOSE_SRC
else:
2022-04-26 10:09:44 -05:00
self._init_new_game(player_color=BLACK)
self._state = GameState.THINKING
2022-04-25 14:49:42 -05:00
# 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]
2022-04-25 14:49:42 -05:00
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:
2022-04-30 11:00:14 -05:00
self._dst_idx = 0
2022-04-25 14:49:42 -05:00
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:
2022-04-30 11:00:14 -05:00
src = self._get_sources()[self._src_idx]
dst = self._get_dests()[self._dst_idx]
2022-04-26 10:09:44 -05:00
self._make_move(src + dst)
2022-04-25 14:49:42 -05:00
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