135 lines
4 KiB
Python
135 lines
4 KiB
Python
import pygame
|
|
import re
|
|
|
|
from constants import WIDTH, HEIGHT, ROWS, COLS, SQUARE_SIZE, LIGHT_BROWN, DARK_BROWN, PIECES
|
|
|
|
|
|
class ChessBoard:
|
|
def __init__(self):
|
|
self.reset()
|
|
|
|
### RL env part
|
|
def reset(self):
|
|
"""
|
|
set up the board and the starting player
|
|
"""
|
|
self.board = self.create_board()
|
|
self.current_player = 'white'
|
|
|
|
def step(self,action):
|
|
"""
|
|
action: A string using natural language and/or chess notation
|
|
Returns: observation, reward, done, info
|
|
"""
|
|
self.handle_input(action)
|
|
# self.move_piece()
|
|
observation = self.board
|
|
reward = self.calculate_reward()
|
|
done = self.check_game_over()
|
|
info = {}
|
|
return observation, reward, done, info
|
|
|
|
def render(self, win):
|
|
"""
|
|
render the pygame window for visual feedback
|
|
"""
|
|
self.display(win)
|
|
pygame.display.update()
|
|
|
|
def close(self):
|
|
pygame.quit()
|
|
|
|
def check_game_over(self):
|
|
if self.draw or self.checkmate or self.resign:
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
### Actual game
|
|
def create_board(self):
|
|
board = [['' for _ in range(COLS)] for _ in range(ROWS)]
|
|
|
|
# Assign unique IDs to each piece
|
|
# Pawns
|
|
for i in range(COLS):
|
|
board[1][i] = 'black_pawn'
|
|
board[6][i] = 'white_pawn'
|
|
|
|
# Rooks
|
|
board[0][0] = board[0][7] = 'black_rook'
|
|
board[7][0] = board[7][7] = 'white_rook'
|
|
|
|
# Knights
|
|
board[0][1] = board[0][6] = 'black_knight'
|
|
board[7][1] = board[7][6] = 'white_knight'
|
|
|
|
# Bishops
|
|
board[0][2] = board[0][5] = 'black_bishop'
|
|
board[7][2] = board[7][5] = 'white_bishop'
|
|
|
|
# Queens
|
|
board[0][3] = 'black_queen'
|
|
board[7][3] = 'white_queen'
|
|
|
|
# Kings
|
|
board[0][4] = 'black_king'
|
|
board[7][4] = 'white_king'
|
|
|
|
return board
|
|
|
|
def display(self, win):
|
|
# draw squares
|
|
for row in range(ROWS):
|
|
for col in range(COLS):
|
|
color = LIGHT_BROWN if (row + col) % 2 == 0 else DARK_BROWN
|
|
pygame.draw.rect(win, color, (col * SQUARE_SIZE, row * SQUARE_SIZE, SQUARE_SIZE, SQUARE_SIZE))
|
|
|
|
# draw pieces
|
|
for row in range(ROWS):
|
|
for col in range(COLS):
|
|
piece = self.board[row][col]
|
|
if piece != '':
|
|
base_piece_name = ''.join([i for i in piece])
|
|
piece_image = PIECES[base_piece_name]
|
|
piece_image = pygame.transform.scale(piece_image, (SQUARE_SIZE, SQUARE_SIZE))
|
|
win.blit(piece_image, (col * SQUARE_SIZE, row * SQUARE_SIZE))
|
|
|
|
def handle_input(self, input_text):
|
|
col_map = {'a':0, 'b':1, 'c':2, 'd':3, 'e':4, 'f':5, 'g':6, 'h':7}
|
|
row_map = {'1':7, '2':6, '3':5, '4':4, '5':3, '6':2, '7':1, '8':0}
|
|
|
|
# TODO: Implement logic for actions without squares
|
|
# This includes statements such as
|
|
# "I capture your queen with my rook"
|
|
|
|
# TODO: Implement logic for actions with mix of squares and names
|
|
# Ex: "Queen to e5"
|
|
# Ex: "a2 captures knight"
|
|
|
|
# TODO: Implement logic for actions with only squares
|
|
# Example: "e4 to e5"
|
|
|
|
# TODO: Probably should spread the different modes across functions for clarity
|
|
# smth like 'figure_pieces()' which figures out the two pieces (and if two pieces interact)
|
|
# and smth like 'figure_squares' which figures out the two squares
|
|
|
|
|
|
def move_piece(self, start_square, end_square):
|
|
|
|
start_col, start_row = start_square
|
|
piece = self.board[start_col][start_row]
|
|
|
|
end_col, end_row = end_square
|
|
end_spot = self.board[end_col][end_row]
|
|
if end_spot == '':
|
|
self.board[end_col][end_row] = piece
|
|
elif 'black' in end_spot:
|
|
self.capture(end_spot)
|
|
self.board[end_col][end_row] = piece
|
|
|
|
self.check_promotion()
|
|
|
|
|
|
|
|
|
|
|