190 lines
6.2 KiB
Python
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)
|