Single_Player_Chess/chess.py

from itertools import product
from unicodedata import lookup
from copy import deepcopy

def cross(A, B):
    """Return sum of of every element from A and B."""
    return tuple(a+b for a in A for b in B)

pieces = ["pawn", "knight", "bishop", "rook", "queen", "king"]
colors = ["white", "black"]

ranks   = "87654321"
files   = "abcdefgh"
squares = cross(files, ranks)

pawn_ranks = {"white" : '2',
              "black" : '7'
             }
home_ranks = {"white" : '1',
              "black" : '8'
             }
init_positions = {"pawn"   : cross(files, pawn_ranks.values()),
                  "rook"   : cross("ah",  home_ranks.values()),
                  "knight" : cross("bg",  home_ranks.values()),
                  "bishop" : cross("cf",  home_ranks.values()),
                  "queen"  : cross("d",   home_ranks.values()),
                  "king"   : cross("e",   home_ranks.values())
                 }

AN_names = {'R' : "rook",
            'N' : "knight",
            'B' : "bishop",
            'Q' : "queen",
            'K' : "king"}

def move(sq, v):
    """Add vector v to square and return a square, even if not in board."""
    return chr(ord(sq[0]) + v[0]) + str(int(sq[1]) + v[1])

def moves(sq, dirs):
    """Add multiple vectors to a square. Return squares that are in board."""
    targets = [move(sq, dir) for dir in dirs]
    return [sq for sq in targets if sq in squares]

def invert(color):
    """Return inverted color."""
    if color == "white":
        return "black"
    return "white"

class Empty:
    """Empty class to avoid unnecessary quotations."""
    def __repr__(self):
        return str(vars(self))

def parse_AN(move):
    """Parse chess move writen in algebraic notation. Return relevant
       information encapsulated within an Empty class, or False if parsing
       fails."""
    def suffix_in(l, dval=False, i=1):
        nonlocal move
        x = move[-i:]
        if x in list(l):
            move = move[:-i]
            return x
        return dval

    queenside = ["{}-{}-{}".format(c,c,c) for c in "0O"]
    kingside  = ["{}-{}".format(c,c)      for c in "0O"]

    r = Empty()
    r.kingside = r.queenside = False
    if move in queenside:
        r.queenside = True
        return r
    if move in kingside:
        r.kingside = True
        return r

    r.move      = move
    r.mate      = suffix_in('#') == '#'
    r.check     = suffix_in('+') == '+'
    r.promotion = suffix_in(AN_names, None)
    if r.promotion:
        r.promotion = AN_names[r.promotion]
        _           = suffix_in('=')
    r.dest    = suffix_in(squares, i=2)
    r.capture = suffix_in("xX") != False
    r.rank    = suffix_in(ranks, None)
    r.file    = suffix_in(files, None)
    r.piece   = suffix_in(AN_names, 'pawn')
    if len(move) > 0 or (r.promotion and r.piece != "pawn") or not r.dest:
        return False
    if r.piece != "pawn":
        r.piece = AN_names[r.piece]
    return r

# Names for vectors
up,   down    = (0, 1),  (0, -1)
left, right   = (-1, 0), (1, 0)
rook_dirs     = (up, down, left, right)

upper_left, upper_right = (-1, 1),  (1, 1)
lower_left, lower_right = (-1, -1), (1, -1)
bishop_dirs             = (upper_left, upper_right,
                           lower_left, lower_right)

class Piece:
    """Contains piece color and piece type"""
    def __init__(self, color=None, piece=None):
        """Assign color and piece. If not specified they
           default to None."""
        self.color = color
        self.piece = piece

    def __repr__(self):
        """Returns unicode representation of instance."""
        if None in (self.piece, self.color):
            return ""
        name = self.color.upper() + " CHESS " + self.piece.upper()
        return lookup(name)

    def __str__(self):
        """Same as repr."""
        return repr(self)

    def __eq__(self, other):
        """Ensures equality testing. If color or piece or both
           equal None, instance equals None."""
        if other == None:
            return  None in (self.color, self.piece)
        if not isinstance(other, Piece):
            return NotImplemented
        return self.piece == other.piece and self.color == other.color

class Game:
    """Chess game class.
       Squares are represented by strings made of letters and integers e.g. 'h4'.
       Board is represented by a dictionary whose keys are squares
       and values instances of Piece class.
       Moves are represented internally by tuples of squares;
       a source and destination.
       Moves are appended to moves stack.
       Upon each move, the relevant attributes of the instance (the game state)
       are appended to the history stack. This allows loading previous
       game states to instance variables."""
    def __init__(self):
        self.index = 0
        self.make_board()
        self.moves = []
        self.AN_moves = []
        self.history = []
        self.turn  = "white"
        self.can_castle = {"white" : {"queen" : True,
                                      "king"  : True},
                           "black" : {"queen" : True,
                                      "king"  : True}
                          }
        self.history.append((deepcopy(self.can_castle),
                             deepcopy(self.board),
                             self.turn))

    def __repr__(self):
        """Return unicode board representation."""
        r = ""
        for rank in ranks:
            r += rank + " |"
            for file in files:
                sq = file+rank
                r += " "
                if self.is_empty(sq):
                    r += " "
                else:
                    r += repr(self.board[sq])
            r += "\n"
        r += "  +" + "-"*16 + "\n"
        r += " "*4 + " ".join(list(files))
        return r

    def make_board(self):
        """Create board and store it in self.board."""
        self.board = dict((sq, Piece()) for sq in squares)
        # Add pieces
        for piece, positions in init_positions.items():
            for sq in positions:
                self.board[sq].piece = piece
        # Add colors
        for sq in cross(files, "12"):
            self.board[sq].color = "white"
        for sq in cross(files, "78"):
            self.board[sq].color = "black"
        # Castling markers
        self.board["a1"].side = "queen"
        self.board["a8"].side = "queen"
        self.board["h1"].side = "king"
        self.board["h8"].side = "king"

    def rebase(self):
        """Shorten moves and history stacks if move is made
           while past game state is loaded into memory. Essentialy
           overrides future game states."""
        self.moves = self.moves[:self.index+1]
        self.history = self.history[:self.index+1]

    def timetravel(self, i):
        """Load game state from past into instance variables.
           Can be reversed since history and moves stack stay intact."""
        if i >= len(self.history) or i < 0:
            return False
        c, board, turn = self.history[i]
        self.board = board
        self.turn = turn
        self.can_castle = c
        self.index = i

    def prev(self):
        """Load game state from previous turn, if it exists."""
        return self.timetravel(self.index-1)

    def next(self):
        """Load game state from next turn, if it exists."""
        return self.timetravel(self.index+1)

    def last(self):
        """Load game state from last move."""
        return self.timetravel(len(self.history)-1)

    def first(self):
        """Load first recorded game state, an empty board."""
        return self.timetravel(0)

    def is_empty(self, sq):
        """Return True if square is empty, else False."""
        return self.board[sq] == None

    def is_check(self, color=None):
        """Return True if king of active player is in check else False."""
        if color == None:
            color = self.turn
        sq = self.occupying(Piece(color, "king"))[0]
        return self.is_attacked(color, sq)

    def is_mate(self):
        """Return True if checkmate for active player."""
        return self.is_check() and self.is_stall()

    def is_stall(self):
        """Return True if active player has no legal move. Whether
           the king is in check is irrelevant."""
        for move in self.all_legal_moves():
            new = deepcopy(self)
            if self.board[move[0]].piece == "pawn" and move[1][1]==home_ranks[invert(self.turn)]:
                move += ("queen",)
            if new.move(*move):
                return False
        return True

    def endofgame(self):
        """Return True if active player can't make a legal move."""
        return self.is_stall() or self.is_mate()

    def occupying(self, piece):
        """Return squares occupied by piece. Piece variable is an
           instance of Piece class."""
        return [sq for sq in squares if self.board[sq] == piece]

    def castle(self, side):
        """Castle if possible, if not return False."""
        color = self.turn
        rank = home_ranks[self.turn]

        if side not in ("queen", "king"):
            return False
        if not self.can_castle[color][side]:
            return False

        if side == "queen":
            check_if_empty    = cross("bcd", rank)
            check_if_attacked = cross("cde", rank)
            rook_move = "a"+rank, "d"+rank
            king_move = "e"+rank, "c"+rank
        elif side == "king":
            check_if_empty    = cross("fg", rank)
            check_if_attacked = check_if_empty
            rook_move = "h"+rank, "f"+rank
            king_move = "e"+rank, "g"+rank
        are_empty    = all(map(self.is_empty, check_if_empty))
        not_attacked = not any(self.is_attacked(color, sq)
                               for sq in check_if_attacked)

        if not (are_empty and not_attacked):
            return False

        self.rebase()
        self.board[king_move[1]] = self.board[king_move[0]]
        self.board[rook_move[1]] = self.board[rook_move[0]]
        self.board[king_move[0]] = Piece()
        self.board[rook_move[0]] = Piece()
        self.can_castle[color]["queen"] = False
        self.can_castle[color]["king"]  = False
        self.moves.append(king_move)
        self.turn = invert(self.turn)
        self.history.append((deepcopy(self.can_castle),
                             deepcopy(self.board),
                             self.turn))
        self.index += 1
        return True

    def move(self, source, target, promotion=None):
        """Move piece from source to target if legal and promote it
           to promotion if it's a pawn whos reached the last
           rank. Return False if fails."""
        new = deepcopy(self)
        if not new.is_legal(source, target, promotion):
            return False

        board = new.board
        moved = board[source]
        eaten = board[target]

        if promotion == None:
            promotion = moved
        else:
            promotion = Piece(new.turn, promotion)

        if moved.piece == "pawn" and target in moves(source, bishop_dirs) and eaten==None:
            # En passant
            a, b = move(target, up), move(target, down)
            board[a] = Piece()
            board[b] = Piece()

        # Castling
        if moved.piece == "rook":
            new.can_castle[new.turn][moved.side] = False
        if eaten.piece == "rook":
            new.can_castle[eaten.color][eaten.side] = False
        if moved.piece == "king":
            new.can_castle[new.turn]["queen"] = False
            new.can_castle[new.turn]["king"]  = False

        board[source] = Piece()
        board[target] = promotion

        if new.is_check(self.turn):
            return False

        self.rebase()
        self.moves.append((source, target))
        self.board = new.board
        self.can_castle = new.can_castle
        self.turn = invert(moved.color)
        self.history.append((deepcopy(self.can_castle),
                             deepcopy(self.board),
                             self.turn))
        self.index += 1
        return True

    def is_legal(self, source, target, promotion):
        """Return True if move is partially legal. Does not
           check whether the active players king is in check
           after the move! See the validate_move method."""
        piece = self.board[source].piece
        color = self.board[source].color
        pawn_on_last_rank = piece == "pawn" and \
                            home_ranks[invert(color)] == target[1]
        if promotion != None and not pawn_on_last_rank:
            return False
        if promotion == None and pawn_on_last_rank:
            return False
        if promotion in ("king", "pawn"):
            return False
        return color == self.turn and target in self.possible_moves(source)

    def is_attacked(self, color, sq):
        """Return True if square is attacked by enemy of opposite color."""
        enemy_color = invert(color)
        attacked = []
        for piece in pieces:
            occupied = self.occupying(Piece(enemy_color, piece))
            attacked += sum(list(map(self.attacks, occupied)), [])
        return sq in attacked

    def all_legal_moves(self):
        """Return all possible partially legal moves by active
           player. These moves can put the active player in check,
           thus they are not completely legal. Castlings not included."""
        color = self.turn
        sqrs = []
        for piece in pieces:
            piece = Piece(color, piece)
            sqrs += self.occupying(piece)
        r = []
        for sq in sqrs:
            mvs = self.possible_moves(sq)
            r += [(sq, m) for m in mvs]
        return r

    def possible_moves(self, sq):
        """Return all squares the piece from sq can move to.
           Moving to these squares might put the player in check.
           Castlings not included."""
        board = self.board
        piece = board[sq].piece
        color = board[sq].color

        def can_eat(target):
            eaten = board[target]
            return eaten != None and eaten.color != color

        is_empty = self.is_empty

        targets = [t for t in self.attacks(sq)
                         if can_eat(t) or is_empty(t)]

        if piece == "pawn":
            r = []
            dir, back = up, down
            if color == "black":
                dir, back = down, up
            frwd = move(sq, dir)
            jump = move(frwd, dir)

            if frwd in squares and is_empty(frwd):
                r.append(frwd)
                is_on_pawn_rank = pawn_ranks[color] == sq[1]
                if is_on_pawn_rank and is_empty(jump):
                    r.append(jump)
            for t in targets:
                a, b = move(t, dir), move(t, back)
                en_passant = can_eat(b) and self.index>0 and\
                             self.moves[self.index-1] == (a,b)
                if can_eat(t) or en_passant:
                    r.append(t)
            return r
        elif piece == "king":
            return [sq for sq in targets if not self.is_attacked(color, sq)]
        else:
            return targets

    def attacks(self, sq):
        """Return all squares the piece on sq attacks. Return empty list
           if square is empty."""
        if sq not in squares:
            return []

        board = self.board
        piece = board[sq].piece
        color = board[sq].color

        def possible_line(dir):
            target = move(sq, dir)
            while target in board and self.is_empty(target):
                yield target
                target = move(target, dir)
            if target in board:
                yield target

        def possible_lines(dirs):
            return sum(list(map(list, map(possible_line, dirs))), [])

        rook   = possible_lines(rook_dirs)
        bishop = possible_lines(bishop_dirs)

        if piece == "pawn":
            if color == "white":
                dirs = upper_left, upper_right
            else:
                dirs = lower_left, lower_right
            return moves(sq, dirs)
        elif piece == "knight":
            nums = [2, -2, 1, -1]
            dirs = [(a, b) for a,b in product(nums, nums)
                               if abs(a) != abs(b)]
            return moves(sq, dirs)
        elif piece == "rook":
            return rook
        elif piece == "bishop":
            return bishop
        elif piece == "queen":
            return rook + bishop
        elif piece == "king":
            dirs = rook_dirs + bishop_dirs
            return moves(sq, dirs)
        return []

    def validate_move(self, source, dest, promotion):
        """Create a copy of the game and try given move.
           Return True if the move succeeds. This method
           exists because is_legal does not test whether
           the move puts the player in check."""
        new = deepcopy(self)
        valid = new.move(source, dest, promotion)
        new.turn = self.turn
        return valid, new

    def AN_move(self, move):
        """Move is given in algebraic notation. The method
           moves the piece if the move is legal and unambiguous,
           else it returns False."""
        AN = move
        info = parse_AN(move)
        if info == False:
            return False
        if info.queenside:
            return self.castle("queen")
        if info.kingside:
            return self.castle("king")

        piece = info.piece
        sqrs = self.occupying(Piece(self.turn, piece))
        promotion = info.promotion
        dest = info.dest
        if info.file != None:
            sqrs = [sq for sq in sqrs if sq[0]==info.file]
        if info.rank != None:
            sqrs = [sq for sq in sqrs if sq[1]==info.rank]
        moves = []
        for sq in sqrs:
            valid, new = self.validate_move(sq, dest, promotion)
            check = new.is_check()==info.check and new.is_mate()==info.mate
            if valid and check:
                moves.append((sq, dest, promotion))
        if len(moves) != 1:
            return False
        move = moves[0]
        self.move(*move)
        self.AN_moves.append(AN)
        return True