moved chess logic to class

This commit is contained in:
Rudis Muiznieks 2022-04-25 14:49:42 -05:00
parent 6620385196
commit ee1ca2f653
Signed by: rudism
GPG Key ID: CABF2F86EF7884F9
3 changed files with 238 additions and 173 deletions

View File

@ -1,176 +1,6 @@
from os import path
from enum import Enum, auto
import time
import re
import itertools
from icecream import ic
from cinput import ControlInput, Button
from cinput import ControlInput
from graphics import Graphics
from .draw import Draw
from . import sunfish
class GameState(Enum):
MAIN_MENU = auto()
THINKING = auto()
CHOOSE_SRC = auto()
CHOOSE_DST = auto()
GAME_OVER = auto()
SAVE_FILE = path.join("data", "chess_state.fen")
MAX_THINKING_SECONDS = 3
from .game import ChessGame
def execute(cinput: ControlInput, graphics: Graphics, _):
graphics.clear()
draw = Draw(graphics)
hist: list[sunfish.Position]
searcher = sunfish.Searcher()
if path.exists(SAVE_FILE):
with open(SAVE_FILE, "r") as fen:
fenstr = fen.read()
hist = [parseFEN(fenstr)]
else:
hist = [sunfish.Position(
sunfish.initial, 0, (True,True), (True,True), 0, 0)]
player_color = 'wb'[get_color(hist[-1])]
draw.draw_board(hist[-1], player_color)
state = GameState.MAIN_MENU
menu_index = 0
move = None
score = None
last_move = "Begin!"
all_moves = dict()
src_idx = 0
dst_idx = 0
src = ""
while True:
key = None
if state == GameState.MAIN_MENU:
draw.draw_menu(menu_index)
key = cinput.get_one_shot()
elif state == GameState.THINKING:
draw.draw_thinking(last_move)
start = time.time()
for _, move, score in searcher.search(hist[-1], hist):
if time.time() - start > MAX_THINKING_SECONDS:
break
hist.append(hist[-1].move(move))
draw.draw_board(hist[-1], player_color)
last_move = move_str(move, player_color == "w")
all_moves = dict()
state = GameState.CHOOSE_SRC
elif state == GameState.CHOOSE_SRC:
if len(list(all_moves)) == 0:
all_moves = get_all_moves(hist[-1])
src = list(all_moves)[src_idx]
draw.draw_select(last_move, list(all_moves)[src_idx])
key = cinput.get_one_shot()
elif state == GameState.CHOOSE_DST:
src = list(all_moves)[src_idx]
draw.draw_select(last_move, src, all_moves[src][dst_idx])
key = cinput.get_one_shot()
elif state == GameState.GAME_OVER:
key = cinput.get_one_shot()
if state == GameState.MAIN_MENU:
if key == Button.DIR_U:
menu_index -= 1
if menu_index < 0:
menu_index = 0
elif key == Button.DIR_D:
menu_index += 1
if menu_index > 3:
menu_index = 3
elif key == Button.BTN_B:
if menu_index == 0:
state = GameState.CHOOSE_SRC
elif menu_index == 1:
with open(SAVE_FILE, "w") as fen:
fen.write(renderFEN(hist[-1]))
return
elif menu_index == 2 or menu_index == 3:
pos = sunfish.Position(
sunfish.initial, 0, (True,True), (True,True), 0, 0)
if menu_index == 3:
player_color = "b"
state = GameState.THINKING
else:
player_color = "w"
state = GameState.CHOOSE_SRC
hist = [pos]
draw.draw_board(hist[-1] if player_color == "w" else hist[-1].rotate(), player_color)
elif state == GameState.CHOOSE_SRC:
if key == Button.DIR_D or key == Button.DIR_R:
src_idx += 1
if src_idx >= len(all_moves.keys()):
src_idx = 0
elif key == Button.DIR_U or key == Button.DIR_L:
src_idx -= 1
if src_idx < 0:
src_idx = len(list(all_moves)) - 1
elif key == Button.BTN_A:
state = GameState.MAIN_MENU
elif key == Button.BTN_B:
state = GameState.CHOOSE_DST
elif state == GameState.CHOOSE_DST:
if key == Button.DIR_D or key == Button.DIR_R:
dst_idx += 1
if dst_idx >= len(all_moves[src]):
dst_idx = 0
elif key == Button.DIR_U or key == Button.DIR_L:
dst_idx -= 1
if dst_idx < 0:
dst_idx = len(all_moves[src]) - 1
elif key == Button.BTN_A:
state = GameState.CHOOSE_SRC
elif key == Button.BTN_B:
dst = all_moves[src][dst_idx]
last_move = src + " - " + dst
sfmove = sunfish.parse(src), sunfish.parse(dst)
hist.append(hist[-1].move(sfmove))
draw.draw_board(hist[-1].rotate(), player_color)
state = GameState.THINKING
else:
if key == Button.BTN_A or key == Button.BTN_B:
state = GameState.MAIN_MENU
def get_all_moves(pos: sunfish.Position):
all_moves = dict()
for src, dst in pos.gen_moves():
all_moves.setdefault(sunfish.render(src), []).append(sunfish.render(dst))
return all_moves
def move_str(move, shift = False):
s = -119 if shift else 0
return sunfish.render(s + move[0]) + " - " + sunfish.render(s + move[1])
def get_color(pos):
return 1 if pos.board.startswith('\n') else 0
def parseFEN(fen):
board, color, castling, enpas, _, _ = fen.split()
board = re.sub(r'\d', (lambda m: '.'*int(m.group(0))), board)
board = list(21*' ' + ' '.join(board.split('/')) + 21*' ')
board[9::10] = ['\n']*12
board = ''.join(board)
wc = ('Q' in castling, 'K' in castling)
bc = ('k' in castling, 'q' in castling)
ep = sunfish.parse(enpas) if enpas != '-' else 0
score = sum(sunfish.pst[p][i] for i,p in enumerate(board) if p.isupper())
score -= sum(sunfish.pst[p.upper()][119-i] for i,p in enumerate(board) if p.islower())
pos = sunfish.Position(board, score, wc, bc, ep, 0)
return pos if color == 'w' else pos.rotate()
def renderFEN(pos, half_move_clock=0, full_move_clock=1):
color = 'wb'[get_color(pos)]
if color == 'b':
pos = pos.rotate()
board = '/'.join(pos.board.split())
board = re.sub(r'\.+', (lambda m: str(len(m.group(0)))), board)
castling = ''.join(itertools.compress('KQkq', pos.wc[::-1]+pos.bc)) or '-'
ep = sunfish.render(pos.ep) if not pos.board[pos.ep].isspace() else '-'
clock = '{} {}'.format(half_move_clock, full_move_clock)
return ' '.join((board, color, castling, ep, clock))
ChessGame(cinput, graphics).run()

194
plugin/chess/game.py Normal file
View File

@ -0,0 +1,194 @@
from os import path
from enum import Enum, auto
import time
from cinput import ControlInput, Button
from graphics import Graphics
from .draw import Draw
from . import sunfish
from . import util
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._cinput = cinput
self._draw = Draw(graphics)
self._state = GameState.MAIN_MENU
self._menu_index = 0
self._searcher = sunfish.Searcher()
def _init_new_game(self, pos: sunfish.Position):
self._hist = [pos]
self._last_move = "Begin!"
self._all_moves = dict()
self._src_idx = 0
self._dst_idx = 0
self._draw.draw_board(self._hist[-1], self._player_color)
def _load_saved_or_init(self):
if path.exists(self.SAVE_FILE):
with open(self.SAVE_FILE, "r") as fen:
fenstr = fen.read()
pos = util.parseFEN(fenstr)
self._player_color = util.get_color(pos)
self._init_new_game(pos)
else:
self._player_color = "w"
self._init_new_game(sunfish.Position(
sunfish.initial, 0, (True,True), (True,True), 0, 0))
def _save_game(self):
with open(self.SAVE_FILE, "w") as fen:
fen.write(util.renderFEN(self._hist[-1]))
def run(self):
# either load the save game or start a new one
self._load_saved_or_init()
while True:
key = None
move = None
src = ""
################
# 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)
start = time.time()
for _, self._move, self._score in self._searcher.search(
self._hist[-1], self._hist):
if time.time() - start > self.MAX_THINKING_SECONDS:
break
self._hist.append(self._hist[-1].move(self._move))
self._draw.draw_board(self._hist[-1], self._player_color)
self._last_move = util.move_str(move, self._player_color == "w")
# reset user moves and move state to user's turn
self._all_moves = dict()
self._state = GameState.CHOOSE_SRC
# user picks source piece
elif self._state == GameState.CHOOSE_SRC:
if len(list(self._all_moves)) == 0:
self._all_moves = util.get_all_moves(self._hist[-1])
src = list(self._all_moves)[self._src_idx]
self._draw.draw_select(self._last_move, src)
key = self._cinput.get_one_shot()
# 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
pos = sunfish.Position(
sunfish.initial, 0, (True,True), (True,True), 0, 0)
if self._menu_index == 3:
self._player_color = "b"
self._state = GameState.THINKING
else:
self._player_color = "w"
self._state = GameState.CHOOSE_SRC
self._init_new_game(pos)
# TODO: figure out why I need to do this
self._draw.draw_board(
pos if self._player_color == "w" else pos.rotate(), self._player_color)
# 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]
# TODO: fix rotation depending on player color
self._last_move = src + " - " + dst
sfmove = sunfish.parse(src), sunfish.parse(dst)
self._hist.append(self._hist[-1].move(sfmove))
self._draw.draw_board(self._hist[-1].rotate(), self._player_color)
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

41
plugin/chess/util.py Normal file
View File

@ -0,0 +1,41 @@
import re
import itertools
from . import sunfish
def get_all_moves(pos: sunfish.Position):
all_moves = dict()
for src, dst in pos.gen_moves():
all_moves.setdefault(sunfish.render(src), []).append(sunfish.render(dst))
return all_moves
def move_str(move, shift = False):
s = -119 if shift else 0
return sunfish.render(s + move[0]) + " - " + sunfish.render(s + move[1])
def get_color(pos):
return "wb"[1 if pos.board.startswith("\n") else 0]
def parseFEN(fen):
board, color, castling, enpas, _, _ = fen.split()
board = re.sub(r"\d", (lambda m: "."*int(m.group(0))), board)
board = list(21*" " + " ".join(board.split("/")) + 21*" ")
board[9::10] = ["\n"]*12
board = "".join(board)
wc = ("Q" in castling, "K" in castling)
bc = ("k" in castling, "q" in castling)
ep = sunfish.parse(enpas) if enpas != "-" else 0
score = sum(sunfish.pst[p][i] for i,p in enumerate(board) if p.isupper())
score -= sum(sunfish.pst[p.upper()][119-i] for i,p in enumerate(board) if p.islower())
pos = sunfish.Position(board, score, wc, bc, ep, 0)
return pos if color == "w" else pos.rotate()
def renderFEN(pos, half_move_clock=0, full_move_clock=1):
color = get_color(pos)
if color == "b":
pos = pos.rotate()
board = "/".join(pos.board.split())
board = re.sub(r"\.+", (lambda m: str(len(m.group(0)))), board)
castling = "".join(itertools.compress("KQkq", pos.wc[::-1]+pos.bc)) or "-"
ep = sunfish.render(pos.ep) if not pos.board[pos.ep].isspace() else "-"
clock = "{} {}".format(half_move_clock, full_move_clock)
return " ".join((board, color, castling, ep, clock))