122 lines
4.3 KiB
Python
122 lines
4.3 KiB
Python
from os import path
|
|
from enum import Enum, auto
|
|
import time
|
|
import re
|
|
import itertools
|
|
from cinput import ControlInput, Button
|
|
from graphics import Graphics
|
|
from .draw import Draw
|
|
from . import sunfish
|
|
|
|
class GameState(Enum):
|
|
MAIN_MENU = auto()
|
|
THINKING = auto()
|
|
CHOOSE_ROW = auto()
|
|
CHOOSE_COL = auto()
|
|
GAME_OVER = auto()
|
|
|
|
SAVE_FILE = path.join("data", "chess_state.fen")
|
|
MAX_THINKING_SECONDS = 3
|
|
|
|
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].board if player_color == 'w' else hist[-1].rotate().board)
|
|
|
|
state = GameState.MAIN_MENU
|
|
menu_index = 0
|
|
move = None
|
|
score = None
|
|
last_move = "Begin!"
|
|
|
|
while True:
|
|
if state == GameState.MAIN_MENU:
|
|
draw.draw_menu(menu_index)
|
|
elif state == GameState.THINKING:
|
|
draw.draw_thinking()
|
|
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].board if player_color == 'b' else hist[-1].rotate().board)
|
|
shift = -119 if player_color == 'b' else 0
|
|
last_move = sunfish.render(shift+move[0]) + " -> " + sunfish.render(shift+move[1])
|
|
if score == sunfish.MATE_UPPER:
|
|
draw.draw_checkmate(last_move, False)
|
|
state = GameState.GAME_OVER
|
|
else:
|
|
draw.draw_rowsel(last_move)
|
|
state = GameState.CHOOSE_ROW
|
|
|
|
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_ROW
|
|
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:
|
|
state = GameState.THINKING
|
|
else:
|
|
state = GameState.CHOOSE_ROW
|
|
hist = [pos]
|
|
draw.draw_board(pos.board)
|
|
|
|
else:
|
|
if key == Button.BTN_A:
|
|
state = GameState.MAIN_MENU
|
|
|
|
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))
|