cheapchess.py

A generator-based, really really dumb random chess move generator

import random
import logging
logger = logging.getLogger("chess")
logger.addHandler(logging.StreamHandler())
import time

import itertools

# the board is a sparse dictionary of x-y addresses
# it contains a color key (W or B) and a generator
# funcio
board = {}
controlled = {'W': set(), 'B': set()}

# these functions generate a series of possible moves from a given
# x, y location.  The main function returns set of generator functions
# each of which represents a series of moves along a possible vector

# for example a bishop could go  X+1, Y+1 -> X+2, Y+2 -> X+3, Y+3...  etc
# OR it could  go on one of the other diagonals; each set of moves
# is represented by a generator function


def bishop(x, y):

    home = x, y

    yield home
    for offset in range(1, 8):
        yield (x + offset, y + offset)

    yield home
    for offset in range(1, 8):
        yield (x - offset, y + offset)

    yield home
    for offset in range(1, 8):
        yield (x + offset, y - offset)

    yield home
    for offset in range(1, 8):
        yield (x - offset, y - offset)


def rook(x, y):

    home = x, y

    yield home
    for yy in range(1, 8):
        yield x, y + yy

    yield home
    for yy in range(1, 8):
        yield x, y - yy

    yield home
    for xx in range(1, 8):
        yield x + xx, y

    yield home
    for xx in range(1, 8):
        yield x - xx, y


# it's easy to make a queen by combining what a bishop
# does and what a rook does

def queen(x, y):
    # inherit 'home' from bishop and rook calls
    for item in bishop(x, y):
        yield item

    for item in rook(x, y):
        yield item

#  A knight or king has only one move set since they don't pass
#  through multiple spaces:


def knight(x, y):

    yield x, y
    yield (x - 1, y + 2)
    yield (x - 2, y + 1)
    yield (x - 1, y - 2)
    yield (x - 2, y - 1)
    yield (x + 1, y + 2)
    yield (x + 2, y + 1)
    yield (x + 1, y - 2)
    yield (x + 2, y - 1)


def king(x, y):
    # for the king, we don't want to move into
    # a space that an enemy could move into

    yield x, y
    yield (x - 1, y + 1)
    yield (x - 1, y)
    yield (x - 1, y - 1)
    yield (x + 1, y + 1)
    yield (x + 1, y)
    yield (x + 1, y - 1)
    yield (x, y - 1)
    yield (x, y + 1)


# a pawn is a bit special, since white pawns and black pawns
# can only do some moves based on captures and cannot backpedal
# to save on repetition, we can genericise it

def pawn(x, y, direction=1, opener=False, enemy_color='W'):

    home = x, y
    yield home
    if board.get((x, y + direction)) is None:
        yield (x, y + direction)
    if opener and board.get((x, y + direction + direction)) is None:
        yield (x, y + direction + direction)

    yield home
    diag = board.get((x + 1, y + direction))
    if diag and diag[0] == enemy_color:
        yield (x + 1, y + direction)

    yield home
    diag = board.get((x - 1, y + direction))
    if diag and diag[0] == enemy_color:
        yield (x - 1, y + direction)


# then we specialize the pawn functions for the two sides
def pawn_w(x, y):
    return pawn(x, y, 1, y == 1)


def pawn_b(x, y):
    return pawn(x, y, -1, y == 6)


# all addresses get clipped from 0-7 in both axes
def clip(addr):
    x, y = addr
    return -1 < x < 8 and -1 < y < 8

    # and when we're sliding along we need to know if we hit a
    # friendly or enemy piece.  We pass in our own color
    # and stop when we hit anything -- but if the color we
    # hit is an enemy it's still a valid move


def slide(addresses, our_color):

    clipped = (m for m  in addresses if clip(m))

    stopped = False
    home = None
    for addr in clipped:
        if home is None:
            home = addr
        if addr == home:
            stopped = False
            continue
        if not stopped:
            next_square = board.get(addr)
            stopped = next_square is not None
            if not stopped or next_square[0] != our_color:
                yield addr



def pick_move(x, y):
    """
    given a board address, check all of the possible moves
    that the piece in that address could make.
    """
    #logger.info('checking (%d,%d)', x, y)

    this_color, move_set = board[x, y]

    all_moves = move_set(x, y)
    sequences = slide(all_moves, this_color)
    possible_moves = tuple(sequences)

    if not possible_moves:
        logger.debug("no move for piece at (%d, %d)", x, y)
        return None

    return random.choice(possible_moves)


def move(W_or_B, started_checked=False):

    friendly_pieces = tuple(
        (addr for addr in board if board[addr][0] == W_or_B))

    if started_checked:
        king_addr = None
        for addr, piece in board.items():
            if piece[1] == king and piece[0] == W_or_B:
                king_addr = addr

        # pad out the list of pieces 
        # so it's more likely a checked king will
        # try to move
        friendly_pieces = [king_addr] * \
            len(friendly_pieces) + list(friendly_pieces)

    # logger.info(friendly_pieces)
    destination = None
    moving_piece = None
    while not destination:
        x, y = random.choice(friendly_pieces)
        destination = pick_move(x, y)
        moving_piece = board[x, y][1]

    if not destination:
        logger.critical("could not find move for piece at %d, %d", x, y)
        logger.setLevel(logging.INFO)
        logger.info(board)
        print_board()
        raise ValueError()

    logger.info('move (%d, %d) %s to (%d, %d)', x, y,
                board[x, y][1].func_name, destination[0], destination[1])

    target, function = board.get(destination, (None, None))

    if function == king:
        return W_or_B

    board[destination] = board[(x, y)]

    if moving_piece == pawn_w and destination[1] == 7:
        board[destination] = 'W', queen
        logger.info('queened!')

    if moving_piece == pawn_b and destination[1] == 0:
        board[destination] = 'B', queen
        logger.info('queened!')

    del board[x, y]
    update_controlled('W')
    update_controlled('B')


def print_board():
    for row in reversed(range(8)):
        r = [str(row)]
        for c in range(8):
            space = board.get((c, row))
            if space:
                color, func = space
                func_name = func.func_name
                if func_name == 'king':
                    func_name = "*"
            else:
                color = " "
                func_name = " "

            r.append(color + func_name[0])

        if row % 2 == 0:
            logger.info("\t--\t\t--\t\t--\t\t--\t")
        else:
            logger.info("\t\t--\t\t--\t\t--\t\t--'")
        logger.info("\t".join(r))
        if row % 2 == 0:
            logger.info("\t--\t\t--\t\t--\t\t--\t")
        else:
            logger.info("\t\t--\t\t--\t\t--\t\t--'")

    logger.info("\t0\t1\t2\t3\t4\t5\t6\t7")


def initialize():
    """
    reset the board dictionary
    """
    board.clear()

    for n in range(8):
        board[n, 1] = 'W', pawn_w
        board[n, 6] = 'B', pawn_b

    board[0, 0] = 'W', rook
    board[7, 0] = 'W', rook
    board[0, 7] = 'B', rook
    board[7, 7] = 'B', rook
    board[1, 0] = 'W', knight
    board[6, 0] = 'W', knight
    board[1, 7] = 'B', knight
    board[6, 7] = 'B', knight
    board[2, 0] = 'W', bishop
    board[5, 0] = 'W', bishop
    board[2, 7] = 'B', bishop
    board[5, 7] = 'B', bishop
    board[3, 0] = 'W', queen
    board[4, 0] = 'W', king
    board[4, 7] = 'B', queen
    board[3, 7] = 'B', king


def update_controlled(color):
    control_set = controlled[color]
    control_set.clear()
    enemy_color = 'W' if color == 'B' else 'W'
    for addr, contents in board.items():
        if contents[0] == color:
            func = contents[1]
            if func in (pawn_w, pawn_b):    
                if func == pawn_w:
                    control_set.add((addr[0] - 1, addr[1] + 1))
                    control_set.add((addr[0] + 1, addr[1] + 1))
                else:
                    control_set.add((addr[0] - 1, addr[1] - 1))
                    control_set.add((addr[0] + 1, addr[1] - 1))
            else:
                contiguous_moves = func(*addr)
                for valid_move in slide(contiguous_moves, color):
                    control_set.add(valid_move)


def is_in_check(color):
    enemy_color = 'W' if color == 'B' else 'W'
    for addr, item in board.items():
        if item == (color, king):
            return addr in controlled[enemy_color]


def simulate(plays, display=False):

    if display:
        logger.setLevel(logging.INFO)
    else:
        logger.setLevel(logging.WARNING)

    wins = {'W': 0, 'B': 0}
    turns = 0
    for p in range(plays):
        initialize()
        for n in range(256):  # an arbitrary cutoff; real world record is 270-ish

            player = 'WB'[n % 2]
            other = 'BW' [n % 2]

            logger.info("\nturn %d\n", n)

            start_checked = is_in_check(player)
            if start_checked:
                logger.info("%s starts checked", player)

            if move(player, start_checked):
                logger.info("%s victory", player)
                wins[other] += 1
                break

            if is_in_check(player):
                logger.info('%s mated!', player)
                wins[other] += 1
                break

            print_board()
            turns += 1


    print_board()
    return {'turns': turns, 'white wins': wins['W'], 'Black wins': wins['B']}


t = time.time()
logger.warning(simulate(1, 0))
logger.warning("elapsed: %d", time.time() - t)