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/__init__.py

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))