Chess-AI-Python/chess_tools.py
2023-03-04 18:42:17 +01:00

668 lines
22 KiB
Python

import typing
from typing import Tuple, List, Iterable, Union
import numpy as np
from numpy import ndarray
import time
from array import array
# The chess board is an 8x8 array of 1-character strings. uppercase = white, lowercase = black
#chess_board = np.empty((8, 8), dtype='U1')
def is_in_board(pos: Tuple[int, int]) -> bool:
# Check if a position is in the board.
row, col = pos
return 0 <= row < 8 and 0 <= col < 8
class CharArray2D:
"""Fast 2D array of characters. Works like a numpy array of chr, but is much faster and memory efficient with pypy"""
__slots__ = ('width', 'height', 'array')
def __init__(self, width, height):
self.width = width
self.height = height
self.array = bytearray(width * height)
def __getitem__(self, index):
y, x = index
return chr(self.array[y * self.width + x])
def __setitem__(self, index, value):
y, x = index
self.array[y * self.width + x] = ord(value)
def copy(self):
new_array = CharArray2D(self.width, self.height)
new_array.array = self.array.copy()
return new_array
@classmethod
def from_ndarray(cls, array):
self = cls(array.shape[1], array.shape[0])
for y in range(self.height):
for x in range(self.width):
self[y, x] = array[y, x]
return self
def __str__(self):
return str(self.array)
def init_board():
# Initialize the board with the starting positions of the pieces.
chess_board = np.array([['R', 'N', 'B', 'Q', 'K', 'B', 'N', 'R'],
['P', 'P', 'P', 'P', 'P', 'P', 'P', 'P'],
['.', '.', '.', '.', '.', '.', '.', '.'],
['.', '.', '.', '.', '.', '.', '.', '.'],
['.', '.', '.', '.', '.', '.', '.', '.'],
['.', '.', '.', '.', '.', '.', '.', '.'],
['p', 'p', 'p', 'p', 'p', 'p', 'p', 'p'],
['r', 'n', 'b', 'q', 'k', 'b', 'n', 'r']])
chess_board = CharArray2D.from_ndarray(np.array(chess_board))
return chess_board
def move_piece(start: Tuple[int, int], end: Tuple[int, int], chess_board: ndarray) -> ndarray:
# Move a piece from start to end.
chess_board = chess_board.copy()
if chess_board[start] not in 'KkPp':
chess_board[end] = chess_board[start]
chess_board[start] = '.'
return chess_board
# Check for pawn
if chess_board[start] == "P":
if end[0] == 7:
chess_board[end] = "Q"
chess_board[start] = '.'
return chess_board
#Check for en passant
if start[1] != end[1] and chess_board[end] == '.':
chess_board[end[0], end[1]] = 'P'
chess_board[start[0], start[1]] = '.'
chess_board[end[0] - 1, end[1]] = '.'
return chess_board
chess_board[end] = chess_board[start]
chess_board[start] = '.'
return chess_board
elif chess_board[start] == "p":
if end[0] == 0:
chess_board[end] = "q"
chess_board[start] = '.'
return chess_board
#Check for en passant
if start[1] != end[1] and chess_board[end] == '.':
chess_board[end[0], end[1]] = 'p'
chess_board[start[0], start[1]] = '.'
chess_board[end[0] + 1, end[1]] = '.'
return chess_board
chess_board[end] = chess_board[start]
chess_board[start] = '.'
return chess_board
# Check for castling.
if chess_board[start] == "K":
if start == (0, 4) and end == (0, 0):
chess_board[0, 2] = "K"
chess_board[0, 3] = "R"
chess_board[0, 0] = '.'
chess_board[0, 1] = '.'
chess_board[0, 4] = '.'
elif start == (0, 4) and end == (0, 7):
chess_board[0, 5] = "R"
chess_board[0, 6] = "K"
chess_board[0, 7] = '.'
chess_board[0, 4] = '.'
else:
chess_board[end] = chess_board[start]
chess_board[start] = '.'
return chess_board
elif chess_board[start] == "k":
if start == (7, 4) and end == (7, 0):
chess_board[7, 2] = "k"
chess_board[7, 3] = "r"
chess_board[7, 0] = '.'
chess_board[7, 1] = '.'
chess_board[7, 4] = '.'
elif start == (7, 4) and end == (7, 7):
chess_board[7, 4] = '.'
chess_board[7, 5] = "r"
chess_board[7, 6] = "k"
chess_board[7, 7] = '.'
else:
chess_board[end] = chess_board[start]
chess_board[start] = '.'
return chess_board
return chess_board
def get_moves_pawn_white(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a white pawn at start.
first_move = (start[0] == 1)
moves = []
# Move forward if not blocked
if chess_board[start[0] + 1, start[1]] == '.':
moves.append((start[0] + 1, start[1]))
if first_move and chess_board[start[0] + 2, start[1]] == '.':
moves.append((start[0] + 2, start[1]))
# Capture diagonally.
if start[1] > 0 and chess_board[start[0] + 1, start[1] - 1].islower():
moves.append((start[0] + 1, start[1] - 1))
if start[1] < 7 and chess_board[start[0] + 1, start[1] + 1].islower():
moves.append((start[0] + 1, start[1] + 1))
# Check for en passant
if start[0] == 4:
if start[1] > 0 and chess_board[start[0], start[1] - 1] == 'p':
moves.append((start[0] + 1, start[1] - 1))
if start[1] < 7 and chess_board[start[0], start[1] + 1] == 'p':
moves.append((start[0] + 1, start[1] + 1))
return moves
def get_moves_pawn_black(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a black pawn at start.
first_move = (start[0] == 6)
moves = []
# Move forward if not blocked.
if chess_board[start[0] - 1, start[1]] == '.':
moves.append((start[0] - 1, start[1]))
if first_move and chess_board[start[0] - 2, start[1]] == '.':
moves.append((start[0] - 2, start[1]))
# Capture diagonally.
if start[1] > 0 and chess_board[start[0] - 1, start[1] - 1].isupper():
moves.append((start[0] - 1, start[1] - 1))
if start[1] < 7 and chess_board[start[0] - 1, start[1] + 1].isupper():
moves.append((start[0] - 1, start[1] + 1))
# Check for en passant
if start[0] == 3:
if start[1] > 0 and chess_board[start[0], start[1] - 1] == 'P':
moves.append((start[0] - 1, start[1] - 1))
if start[1] < 7 and chess_board[start[0], start[1] + 1] == 'P':
moves.append((start[0] - 1, start[1] + 1))
return moves
def get_moves_rook_white(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a white rook at start.
moves = []
# Move up.
for i in range(start[0] + 1, 8):
if chess_board[i, start[1]] == '.':
moves.append((i, start[1]))
elif chess_board[i, start[1]].islower():
moves.append((i, start[1]))
break
else:
break
# Move down.
for i in range(start[0] - 1, -1, -1):
if chess_board[i, start[1]] == '.':
moves.append((i, start[1]))
elif chess_board[i, start[1]].islower():
moves.append((i, start[1]))
break
else:
break
# Move right.
for i in range(start[1] + 1, 8):
if chess_board[start[0], i] == '.':
moves.append((start[0], i))
elif chess_board[start[0], i].islower():
moves.append((start[0], i))
break
else:
break
# Move left.
for i in range(start[1] - 1, -1, -1):
if chess_board[start[0], i] == '.':
moves.append((start[0], i))
elif chess_board[start[0], i].islower():
moves.append((start[0], i))
break
else:
break
return moves
def get_moves_rook_black(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a black rook at start.
moves = []
# Move up.
for i in range(start[0] + 1, 8):
if chess_board[i, start[1]] == '.':
moves.append((i, start[1]))
elif chess_board[i, start[1]].isupper():
moves.append((i, start[1]))
break
else:
break
# Move down.
for i in range(start[0] - 1, -1, -1):
if chess_board[i, start[1]] == '.':
moves.append((i, start[1]))
elif chess_board[i, start[1]].isupper():
moves.append((i, start[1]))
break
else:
break
# Move right.
for i in range(start[1] + 1, 8):
if chess_board[start[0], i] == '.':
moves.append((start[0], i))
elif chess_board[start[0], i].isupper():
moves.append((start[0], i))
break
else:
break
# Move left.
for i in range(start[1] - 1, -1, -1):
if chess_board[start[0], i] == '.':
moves.append((start[0], i))
elif chess_board[start[0], i].isupper():
moves.append((start[0], i))
break
else:
break
return moves
def get_moves_knight_white(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a white knight at start.
moves = []
verify = lambda move: is_in_board(move) and (chess_board[move].islower() or chess_board[move] == '.')
# Move up right.
move = (start[0] + 2, start[1] + 1)
if verify(move):
moves.append(move)
# Move up left.
move = (start[0] + 2, start[1] - 1)
if verify(move):
moves.append(move)
# Move down right.
move = (start[0] - 2, start[1] + 1)
if verify(move):
moves.append(move)
# Move down left.
move = (start[0] - 2, start[1] - 1)
if verify(move):
moves.append(move)
# Move right up.
move = (start[0] + 1, start[1] + 2)
if verify(move):
moves.append(move)
# Move right down.
move = (start[0] - 1, start[1] + 2)
if verify(move):
moves.append(move)
# Move left up.
move = (start[0] + 1, start[1] - 2)
if verify(move):
moves.append(move)
# Move left down.
move = (start[0] - 1, start[1] - 2)
if verify(move):
moves.append(move)
return moves
def get_moves_knight_black(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a black knight at start.
moves = []
verify = lambda move: is_in_board(move) and (chess_board[move].isupper() or chess_board[move] == '.')
# Move up right.
move = (start[0] + 2, start[1] + 1)
if verify(move):
moves.append(move)
# Move up left.
move = (start[0] + 2, start[1] - 1)
if verify(move):
moves.append(move)
# Move down right.
move = (start[0] - 2, start[1] + 1)
if verify(move):
moves.append(move)
# Move down left.
move = (start[0] - 2, start[1] - 1)
if verify(move):
moves.append(move)
# Move right up.
move = (start[0] + 1, start[1] + 2)
if verify(move):
moves.append(move)
# Move right down.
move = (start[0] - 1, start[1] + 2)
if verify(move):
moves.append(move)
# Move left up.
move = (start[0] + 1, start[1] - 2)
if verify(move):
moves.append(move)
# Move left down.
move = (start[0] - 1, start[1] - 2)
if verify(move):
moves.append(move)
return moves
def get_moves_bishop_white(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a white bishop at start.
moves = []
# Move up and right.
for i in range(1, 8):
if start[0] + i > 7 or start[1] + i > 7:
break
if chess_board[start[0] + i, start[1] + i] == '.':
moves.append((start[0] + i, start[1] + i))
elif chess_board[start[0] + i, start[1] + i].islower():
moves.append((start[0] + i, start[1] + i))
break
else:
break
# Move up and left.
for i in range(1, 8):
if start[0] + i > 7 or start[1] - i < 0:
break
if chess_board[start[0] + i, start[1] - i] == '.':
moves.append((start[0] + i, start[1] - i))
elif chess_board[start[0] + i, start[1] - i].islower():
moves.append((start[0] + i, start[1] - i))
break
else:
break
# Move down and right.
for i in range(1, 8):
if start[0] - i < 0 or start[1] + i > 7:
break
if chess_board[start[0] - i, start[1] + i] == '.':
moves.append((start[0] - i, start[1] + i))
elif chess_board[start[0] - i, start[1] + i].islower():
moves.append((start[0] - i, start[1] + i))
break
else:
break
# Move down and left.
for i in range(1, 8):
if start[0] - i < 0 or start[1] - i < 0:
break
if chess_board[start[0] - i, start[1] - i] == '.':
moves.append((start[0] - i, start[1] - i))
elif chess_board[start[0] - i, start[1] - i].islower():
moves.append((start[0] - i, start[1] - i))
break
else:
break
return moves
def get_moves_bishop_black(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a black bishop at start.
moves = []
# Move up and right.
for i in range(1, 8):
if start[0] + i > 7 or start[1] + i > 7:
break
if chess_board[start[0] + i, start[1] + i] == '.':
moves.append((start[0] + i, start[1] + i))
elif chess_board[start[0] + i, start[1] + i].isupper():
moves.append((start[0] + i, start[1] + i))
break
else:
break
# Move up and left.
for i in range(1, 8):
if start[0] + i > 7 or start[1] - i < 0:
break
if chess_board[start[0] + i, start[1] - i] == '.':
moves.append((start[0] + i, start[1] - i))
elif chess_board[start[0] + i, start[1] - i].isupper():
moves.append((start[0] + i, start[1] - i))
break
else:
break
# Move down and right.
for i in range(1, 8):
if start[0] - i < 0 or start[1] + i > 7:
break
if chess_board[start[0] - i, start[1] + i] == '.':
moves.append((start[0] - i, start[1] + i))
elif chess_board[start[0] - i, start[1] + i].isupper():
moves.append((start[0] - i, start[1] + i))
break
else:
break
# Move down and left.
for i in range(1, 8):
if start[0] - i < 0 or start[1] - i < 0:
break
if chess_board[start[0] - i, start[1] - i] == '.':
moves.append((start[0] - i, start[1] - i))
elif chess_board[start[0] - i, start[1] - i].isupper():
moves.append((start[0] - i, start[1] - i))
break
else:
break
return moves
def get_moves_queen_white(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a white queen at start.
return get_moves_bishop_white(start, chess_board) + get_moves_rook_white(start, chess_board)
def get_moves_queen_black(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a black queen at start.
return get_moves_bishop_black(start, chess_board) + get_moves_rook_black(start, chess_board)
def get_moves_king_white(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a black king at start.
moves = []
# Move up if possible.
if start[0] < 7:
if chess_board[start[0] + 1, start[1]] == '.' or chess_board[start[0] + 1, start[1]].islower():
moves.append((start[0] + 1, start[1]))
# Move down.
if start[0] > 0:
if chess_board[start[0] - 1, start[1]] == '.' or chess_board[start[0] - 1, start[1]].islower():
moves.append((start[0] - 1, start[1]))
# Move right.
if start[1] < 7:
if chess_board[start[0], start[1] + 1] == '.' or chess_board[start[0], start[1] + 1].islower():
moves.append((start[0], start[1] + 1))
# Move left.
if start[1] > 0:
if chess_board[start[0], start[1] - 1] == '.' or chess_board[start[0], start[1] - 1].islower():
moves.append((start[0], start[1] - 1))
# Move up and right.
if start[0] < 7:
if start[1] < 7:
if chess_board[start[0] + 1, start[1] + 1] == '.' or chess_board[start[0] + 1, start[1] + 1].islower():
moves.append((start[0] + 1, start[1] + 1))
# Move up and left.
if start[0] < 7:
if start[1] > 0:
if chess_board[start[0] + 1, start[1] - 1] == '.' or chess_board[start[0] + 1, start[1] - 1].islower():
moves.append((start[0] + 1, start[1] - 1))
# Move down and right.
if start[0] > 0:
if start[1] < 7:
if chess_board[start[0] - 1, start[1] + 1] == '.' or chess_board[start[0] - 1, start[1] + 1].islower():
moves.append((start[0] - 1, start[1] + 1))
# Move down and left.
if start[0] > 0:
if start[1] > 0:
if chess_board[start[0] - 1, start[1] - 1] == '.' or chess_board[start[0] - 1, start[1] - 1].islower():
moves.append((start[0] - 1, start[1] - 1))
# Castling.
if start == (0, 4):
# White king side.
if chess_board[0, 5] == '.' and chess_board[0, 6] == '.' and chess_board[0, 7] == 'R':
moves.append((0, 7))
# White queen side.
if chess_board[0, 3] == '.' and chess_board[0, 2] == '.' and chess_board[0, 1] == '.' and chess_board[0, 0] == 'R':
moves.append((0, 0))
return moves
def get_moves_king_black(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a black king at start.
moves = []
# Move up if possible.
if start[0] < 7:
if chess_board[start[0] + 1, start[1]] == '.' or chess_board[start[0] + 1, start[1]].isupper():
moves.append((start[0] + 1, start[1]))
# Move down.
if start[0] > 0:
if chess_board[start[0] - 1, start[1]] == '.' or chess_board[start[0] - 1, start[1]].isupper():
moves.append((start[0] - 1, start[1]))
# Move right.
if start[1] < 7:
if chess_board[start[0], start[1] + 1] == '.' or chess_board[start[0], start[1] + 1].isupper():
moves.append((start[0], start[1] + 1))
# Move left.
if start[1] > 0:
if chess_board[start[0], start[1] - 1] == '.' or chess_board[start[0], start[1] - 1].isupper():
moves.append((start[0], start[1] - 1))
# Move up and right.
if start[0] < 7:
if start[1] < 7:
if chess_board[start[0] + 1, start[1] + 1] == '.' or chess_board[start[0] + 1, start[1] + 1].isupper():
moves.append((start[0] + 1, start[1] + 1))
# Move up and left.
if start[0] < 7:
if start[1] > 0:
if chess_board[start[0] + 1, start[1] - 1] == '.' or chess_board[start[0] + 1, start[1] - 1].isupper():
moves.append((start[0] + 1, start[1] - 1))
# Move down and right.
if start[0] > 0:
if start[1] < 7:
if chess_board[start[0] - 1, start[1] + 1] == '.' or chess_board[start[0] - 1, start[1] + 1].isupper():
moves.append((start[0] - 1, start[1] + 1))
# Move down and left.
if start[0] > 0:
if start[1] > 0:
if chess_board[start[0] - 1, start[1] - 1] == '.' or chess_board[start[0] - 1, start[1] - 1].isupper():
moves.append((start[0] - 1, start[1] - 1))
# Castling.
if start == (7, 4):
# Castling kingside.
if chess_board[7, 5] == '.' and chess_board[7, 6] == '.' and chess_board[7, 7] == 'r':
moves.append((7, 7))
# Castling queenside.
if chess_board[7, 3] == '.' and chess_board[7, 2] == '.' and chess_board[7, 1] == '.' and chess_board[7, 0] == 'r':
moves.append((7, 0))
return moves
def get_possible_moves(start: Tuple[int, int], chess_board: ndarray) -> List[Tuple[int, int]]:
# Get all possible moves for a piece at start.
piece = chess_board[start]
if piece == 'P':
return get_moves_pawn_white(start, chess_board)
elif piece == 'p':
return get_moves_pawn_black(start, chess_board)
elif piece == 'R':
return get_moves_rook_white(start, chess_board)
elif piece == 'r':
return get_moves_rook_black(start, chess_board)
elif piece == 'N':
return get_moves_knight_white(start, chess_board)
elif piece == 'n':
return get_moves_knight_black(start, chess_board)
elif piece == 'B':
return get_moves_bishop_white(start, chess_board)
elif piece == 'b':
return get_moves_bishop_black(start, chess_board)
elif piece == 'Q':
return get_moves_queen_white(start, chess_board)
elif piece == 'q':
return get_moves_queen_black(start, chess_board)
elif piece == 'K':
return get_moves_king_white(start, chess_board)
elif piece == 'k':
return get_moves_king_black(start, chess_board)
return []
def get_time(func):
def wrapper(*args, **kwargs):
start = time.perf_counter()
result = func(*args, **kwargs)
print(f"{func.__name__} took {time.perf_counter() - start} seconds")
return result
return wrapper