chess-ai-py/chess.py
2022-11-26 05:33:39 -08:00

190 lines
6.2 KiB
Python

#!/bin/python
import random
from collections import namedtuple
from vector import vec2, vec2_from_text
Piece = namedtuple('Piece', 'side rank')
ranks = 'RNBQKBNR'
def next_turn(turn):
return 'B' if turn == 'W' else 'W'
def get_turn_pieces(turn, whites, blacks):
return (whites, blacks) if turn == 'W' else (blacks, whites)
def create_board():
whites = {}
whites.update({vec2(c, 0): rank for (c, rank) in zip(range(8), ranks)})
whites.update({vec2(c, 1): 'P' for c in range(8)})
blacks = {}
blacks.update({vec2(c, 7): rank for (c, rank) in zip(range(8), ranks)})
blacks.update({vec2(c, 6): 'P' for c in range(8)})
return whites, blacks
def print_board(whites, blacks):
for y in range(7, -1, -1):
print(str(y + 1) + ' ', end='')
for x in range(8):
pos = vec2(x, y)
if pos in whites:
print('+' + whites[pos] + '+ ', end='')
elif pos in blacks:
print('-' + blacks[pos] + '- ', end='')
else:
print(' ', end='')
print()
print(' ' + ' '.join('ABCDEFGH'))
def get_human_move(turn):
s = input(turn + ' to move: ')
return tuple(map(vec2_from_text, s.split()))
def get_start_move(start_move):
print(start_move)
return tuple(map(vec2_from_text, start_move.split()))
def get_computer_move(turn, whites, blacks):
return random.choice(get_valid_moves(turn, whites, blacks))
def get_valid_dsts(src, turn, whites, blacks):
# todo: castling
# todo: en passant
# todo: promotion
self_pieces, oppo_pieces = get_turn_pieces(turn, whites, blacks)
piece = self_pieces[src]
valid_dsts = set()
straight_dirs = set([vec2(1, 0), vec2(-1, 0), vec2(0, 1), vec2(0, -1)])
diagonal_dirs = set([vec2(1, 1), vec2(-1, 1), vec2(1, -1), vec2(-1, -1)])
def add_dirs(dirs, *, single=False):
for dir in dirs:
delta = dir
while (src + delta).is_valid():
dst = src + delta
if dst in self_pieces:
break
valid_dsts.add(dst)
if dst in oppo_pieces:
break
if single:
break
delta += dir
if piece == 'P': # Pawn
dy = 1 if turn == 'W' else -1
sy = 1 if turn == 'W' else 6
dst_fw1 = src + vec2(0, dy)
dst_fw2 = src + vec2(0, dy * 2)
dst_tkl = src + vec2(-1, dy)
dst_tkr = src + vec2(1, dy)
if dst_fw1.is_valid() and dst_fw1 not in self_pieces and dst_fw1 not in oppo_pieces:
valid_dsts.add(dst_fw1)
if src.y == sy and dst_fw2 not in self_pieces and dst_fw2 not in oppo_pieces:
valid_dsts.add(dst_fw2)
if dst_tkl in oppo_pieces:
valid_dsts.add(dst_tkl)
if dst_tkr in oppo_pieces:
valid_dsts.add(dst_tkr)
elif piece == 'N': # Knight
knight_dirs = set([vec2(1, 2), vec2(2, 1), vec2(2, -1), vec2(1, -2), vec2(-1, -2), vec2(-2, -1), vec2(-2, 1), vec2(-1, 2)])
add_dirs(knight_dirs, single=True)
elif piece == 'R': # Rook
add_dirs(straight_dirs)
elif piece == 'B': # Bishop
add_dirs(diagonal_dirs)
elif piece == 'Q': # Queen
add_dirs(straight_dirs | diagonal_dirs)
elif piece == 'K': # King
add_dirs(straight_dirs | diagonal_dirs, single=True)
return valid_dsts
def get_valid_moves(turn, whites, blacks):
self_pieces = whites if turn == 'W' else blacks
return [(src, dst) for src in self_pieces for dst in get_valid_dsts(src, turn, whites, blacks)]
def next_board(src, dst, turn, whites, blacks):
self_pieces, oppo_pieces = get_turn_pieces(turn, whites, blacks)
self_pieces[dst] = self_pieces[src] # Move src -> dst
del self_pieces[src] # Remove src
if dst in oppo_pieces: # Remove dst if it was an opponent
del oppo_pieces[dst]
return whites, blacks
# Valid moves filtered such that the opponent cannot take the current player's king
def get_legal_moves(turn, whites, blacks):
self_pieces, oppo_pieces = get_turn_pieces(turn, whites, blacks)
valid_moves = get_valid_moves(turn, whites, blacks)
legal_moves = []
for src, dst in valid_moves:
n_whites, n_blacks = next_board(src, dst, turn, whites.copy(), blacks.copy())
n_turn = next_turn(turn)
n_self_pieces, n_oppo_pieces = get_turn_pieces(n_turn, n_whites, n_blacks)
n_valid_moves = get_valid_moves(n_turn, n_whites, n_blacks)
legal = True
for n_src, n_dst in n_valid_moves:
if n_dst in n_oppo_pieces and n_oppo_pieces[n_dst] == 'K':
legal = False
break
if legal:
legal_moves.append((src, dst))
return legal_moves
def is_in_check(turn, whites, blacks):
self_pieces, oppo_pieces = get_turn_pieces(turn, whites, blacks)
n_turn = next_turn(turn)
threat_moves = get_valid_moves(n_turn, whites, blacks)
for src, dst in threat_moves:
if dst in self_pieces and self_pieces[dst] == 'K':
return True
return False
turn = 'W'
whites, blacks = create_board()
start_moves = [
'D2 D4',
'E7 E5',
'D4 E5',
'D8 E7',
'E5 E6',
'E7 E6',
'E2 E4',
'F7 F5',
#'E4 F5' # Should not be legal
]
while True:
print(turn, 'to move')
print_board(whites, blacks)
legal_moves = get_legal_moves(turn, whites, blacks)
print(len(legal_moves), 'legal moves')
self_pieces, oppo_pieces = get_turn_pieces(turn, whites, blacks)
if len(legal_moves) == 0:
print('game is over')
break
# if turn == 'W':
if len(start_moves):
src, dst = get_start_move(start_moves.pop(0))
else:
for src, dst in legal_moves:
print(f"{self_pieces[src]} {src}{'x' if dst in oppo_pieces else ' '}{dst}")
while True:
try:
src, dst = get_human_move(turn)
if not (src, dst) in legal_moves:
print('not a legal move')
else:
break
except ValueError:
print('invalid input')
# else:
# src, dst = get_computer_move(turn, whites, blacks)
whites, blacks = next_board(src, dst, turn, whites, blacks)
turn = next_turn(turn)