pygo/lib/goban.py

from gomill import sgf

# leftoffat: implementing navigation forwards and backward through the game.

class Goban:
    """Represents the go board. Handles stone placement, captures, etc"""

    # enum values for the board array
    EMPTY='.'
    WHITE='w'
    BLACK='b'
    SCORE_BLACK='B'
    SCORE_WHITE='W'
    SCORE_DAME='d'
    SCORING='s'


    def __init__(self, board_size=19, file_name=None):
        # Build the board intersections
        self.board_size = board_size
        num_points = board_size * board_size
        self.board = [Goban.EMPTY] * num_points

        self.def_draw_codes = self._make_default_draw_codes()

        self.to_move = Goban.BLACK
        self.black_captures = 0
        self.white_captures = 0
        self.last_move = None
        self.passed_last = False
        self.ko = None
        self.hover = None
        self.elapsed_time = 0
        self.winner = Goban.EMPTY
        self.move = 0

        self.file_name = file_name

        # Track our game in an easily saveable format!
        self.sgf_game = None

        if self.file_name is not None:
            self.load_sgf(file_name)
        else:
            self.sgf_game = sgf.Sgf_game(self.board_size)


    def load_sgf(self, file_name):
        try:
            with open(file_name, 'r') as fn:
                self.sgf_game = sgf.Sgf_game.from_string(fn.read())
        except IOError:
            # fixme - this should be convertable into a dialog box... perhaps it should throw an exception of its own
            print 'There was a problem loading the SGF file.' 
        
        # Do initial layout
        self._place_initial_stones()

        for node in self.sgf_game.get_main_sequence():
            self._play_node(node)

        if self.sgf_game.get_winner() is not None:
            self.winner = self.sgf_game.get_winner()
            self.to_move = None


    def _place_initial_stones(self):
        root = self.sgf_game.get_root()
        if root.has_setup_stones():
            black, white, empty = root.get_setup_stones()
            for point in black:
                self.board[self._real_pos(self._sgf_to_move(point))] = Goban.BLACK
            for point in white:
                self.board[self._real_pos(self._sgf_to_move(point))] = Goban.WHITE
            for point in empty:
                self.board[self._real_pos(self._sgf_to_move(point))] = Goban.EMPTY


    # Play a single node from an sgf tree. Useful for warking the tree.
    def _play_node(self, node):
        color, pos = node.get_move()
        if color is not None:
            if pos is None:
                self.pass_move(color)
            else:
                self.play_move(self._sgf_to_move(pos), color, add_sgf=False)
        # fixme - add resignation detection


    # Play all moves up to and including node, which should be
    # an sgf.Tree_node
    def play_to_node(self, node):
        self.soft_reset()
        for n in self.sgf_game.get_main_sequence():
            self._sgf_play_node(self, n)
            if n == node:
                return

    
    def play_to_move_n(self, n=None):
        ''' Play to the nth move. If n is None, play to the latest move '''
        self.soft_reset()

        if n is None or len(self.sgf_game.get_main_sequence()) < n:
            n = len(self.sgf_game.get_main_sequence())
        
        portion = self.sgf_game.get_main_sequence()[:n]
        for node in portion:
            self._play_node(node)


    def save_sgf(self, file_name=None):
        '''Saves the current game as an SGF file. If file_name is None, we use the previously specified filename.
        If there is no previously specified filename, we raise an exception.'''

        if self.file_name is None and file_name is not None:
            self.file_name = file_name

        if file_name is None:
            file_name = self.file_name
        if file_name is None:
            return # fixme - this should be an exception instead

        with open(file_name, 'w') as fn:
            fn.write(self.sgf_game.serialise())


    def set_hover(self, pos):
        rpos = self._real_pos(pos)
        if rpos == self.hover:
            return

        if not self._valid_move(rpos) or self.to_move == Goban.EMPTY:
            self.clear_hover()
            return

        self.hover = rpos

    
    def soft_reset(self):
        """Reset the board to a pre-game state, but preserves move history.
        In other words, goes back to before move #1"""
        # Clear the board by setting it to the same size it currently is at
        self.set_board_size(self.board_size)
        self.to_move = Goban.BLACK
        self.black_captures = 0
        self.white_captures = 0
        self.last_move = None
        self.passed_last = False
        self.ko = None
        self.hover = None
        self.elapsed_time = 0
        self.winner = Goban.EMPTY
        self.move = 0

        self._place_initial_stones()


    def reset(self):
        '''Fully resets the game. The only thing preserved is the SGF filename, if it is set '''
        self.sgf_game = sgf.Sgf_game(self.board_size)
        self.soft_reset()


    def set_board_size(self, new_size):
        """Set the board to a new size. This will also
        reset the board to a blank state, but will *not* reset captures, etc
        (call reset() first if you want that)"""
        self.board_size = new_size
        num_points = self.board_size * self.board_size
        self.board = [Goban.EMPTY] * num_points


    def clear_hover(self):
        self.hover = None
        

    def play_move(self, pos, color=None, add_sgf=True):
        '''
        Plays a move
        
        pos: a tuple containing row and column to play on

        color: which color is playing. If not specified, the player whose turn it is is assumed

        add_sgf: if True, this is a new move and should be recorded in the move history

        return:
        To help with drawing code efficiency, any modified positions are returned in a list.
        This includes anything that the GUI may want to redraw in the wake of this move.
        

        '''
        if add_sgf and self.move != len(self.sgf_game.get_main_sequence()):
            return

        if color is None:
            color = self.to_move

        if self.to_move == Goban.EMPTY:
            return None

        rpos = self._real_pos(pos)
        
        if not self._valid_move(rpos, color):
            return None

        self.board[rpos] = color
        self._changed = []
        self._changed.append(rpos)

        if self.ko is not None:
            self._changed.append(self.ko)
        if self.last_move is not None:
            self._changed.append(self.last_move)

        self._capture(rpos)
        self.last_move = rpos
        self.passed_last = False

        if add_sgf:
            node = self.sgf_game.extend_main_sequence()
            node.set_move(color, self._pos_to_sgf(pos))

        self.to_move = self._other_color(color)
        self.clear_hover()

        # If there is a new ko, send that back too
        if self.ko is not None:
            self._changed.append(self.ko)

        self.move += 1

        return self._changed



    def pass_move(self, color=None):        
        if color is None:
            color = self.to_move

        # If the game is over, fail silently
        if color is None:
            return

        self._changed = []
        if self.ko is not None:
            self._changed.append(self.ko)
        if self.last_move is not None:
            self._changed.append(self.last_move)

        node = self.sgf_game.extend_main_sequence()
        node.set_move(color, None)

        if self.passed_last:
            self.to_move = Goban.EMPTY
            self.winner = Goban.SCORING
            self.auto_score()
            return range(len(self.board))
        else:
            self.to_move = self._other_color(color)
            self.passed_last = True

        self.last_move = None
        self.ko = None

        return self._changed


    def resign(self, color=None):
        if color is None:
            color = self.to_move

        # If the game is over, fail silently
        if color is None:
            return

        self._changed = []
        if self.ko is not None:
            self._changed.append(self.ko)
        if self.last_move is not None:
            self._changed.append(self.last_move)

        self.passed_last = False
        self.last_move = None
        self.ko = None
        self.winner = self._other_color(color)
        self.to_move = Goban.EMPTY

        return self._changed


    def _capture(self, pos, color=None):
        """Look for stones captured on the 4 sides of pos, remove them and increment
        capture counter. This pos must be a *real* position value, not an x,y tuple."""

        if color is None:
            color = self.to_move

        # If we get here, we've definitely played a move,
        # clearing any existing ko point
        self.ko = None

        # Who are we capturing
        who = self._other_color(color)

        captures = 0

        for p in self._neighbors(pos):
            if not self._on_board(p):
                continue
                
            if not self._num_liberties(p, who):
                captures += self._delete_group(p)

        if color == Goban.BLACK:
            self.black_captures += captures
        elif color == Goban.WHITE:
            self.white_captures += captures

        # Check for ko
        if captures == 1 and self._num_liberties(pos, color) == 1:
            # find the empty point
            for p in self._neighbors(pos):
                if self.board[p] == Goban.EMPTY:
                    self.ko = p
                    break



    def _valid_move(self, pos, color=None):
        if not self._on_board(pos):
            return False

        if color is None:
            color = self.to_move

        # Can't play atop another stone or on the ko point
        if self.board[pos] != Goban.EMPTY or pos == self.ko:
            return False

        # Temporarily place the stone
        self.board[pos] = color

        liberties = self._num_liberties(pos, color)
        opponent = self._other_color(color)

        kills_group = False
        for d in self._neighbors(pos):
            if not self._on_board(d):
                continue
            
            if self._num_liberties(d, opponent) == 0:
                kills_group = True
                break

        # Remove temporary stone
        self.board[pos] = Goban.EMPTY

        return liberties > 0 or kills_group



    # Recursively find whether there are liberties for the group
    # at pos.
    def _num_liberties(self, pos, who):
        if not self._on_board(pos) or self.board[pos] != who:
            return -1

        bs = self.board_size * self.board_size
        checked = [False] * bs

        return self._num_liberties_r(pos, who, checked)


    def _num_liberties_r(self, pos, who, checked=None):
        if checked[pos]:
            return 0
        else:
            checked[pos] = True

        square = self.board[pos]

        if square == Goban.EMPTY:
            return 1
        elif square != who:
            return 0
        else:
            liberties = 0
            for d in self._neighbors(pos):
                liberties += self._num_liberties_r(d, who, checked)

            return liberties



    # We don't need to worry about crossing ourselves with the
    # recursion here, because we've already deleted the stone.
    def _delete_group(self, pos):
        if not self._on_board(pos):
            return

        who = self.board[pos]

        if who == Goban.EMPTY:
            return 0

        return self._delete_group_r(pos, who)


    def _delete_group_r(self, pos, who):
        if not self._on_board(pos):
            return 0

        if self.board[pos] != who:
            return 0

        self.board[pos] = Goban.EMPTY
        self._changed.append(pos)

        num_deleted = 1

        num_deleted += self._delete_group_r(pos + 1, who)
        num_deleted += self._delete_group_r(pos - 1, who)
        num_deleted += self._delete_group_r(pos + self.board_size, who)
        num_deleted += self._delete_group_r(pos - self.board_size, who)

        return num_deleted


    def draw_code(self, pos):
        if not self._on_board(pos):
            return None

        point = self.board[pos]
        code = None

        if point == Goban.EMPTY or point == Goban.SCORE_DAME:
            code = self.def_draw_codes[pos]
        elif point == Goban.BLACK:
            code = 'b'
            if pos == self.last_move:
                code += 'Cw'
        elif point == Goban.WHITE:
            code = 'w'
            if pos == self.last_move:
                code += 'Cb'
        elif point == Goban.SCORE_WHITE:
            code = self.def_draw_codes[pos] + 'ws'
        elif point == Goban.SCORE_BLACK:
            code = self.def_draw_codes[pos] + 'bs'
            
        if pos == self.ko:
            code += 'S'

        return code
        

    def auto_score(self):
        '''
        Detects regions and assigns them to the appropriate player
        This may not always guess correctly, so we also have API that
        can manually change these things.
        
        After calling this function, the entire board should be redrawn.
        '''
        for i in range(len(self.board)):
            if self.board[i] == Goban.EMPTY:
                bs = self.board_size * self.board_size
                checked = set()
                score = self._score_space(i, checked)

                for c in checked:
                    self.board[c] = score



    def _score_space(self, pos, checked):
        if pos in checked:
            return None

        if self.board[pos] == Goban.BLACK:
            return Goban.SCORE_BLACK
        elif self.board[pos] == Goban.WHITE:
            return Goban.SCORE_WHITE

        checked.add(pos)

        possible = set()
        for i in self._neighbors(pos):
            score = self._score_space(i, checked)
            possible.add(score)

        if Goban.SCORE_DAME in possible or (Goban.SCORE_BLACK in possible and Goban.SCORE_WHITE in possible):
            return Goban.SCORE_DAME
        elif Goban.SCORE_BLACK in possible:
            return Goban.SCORE_BLACK
        elif Goban.SCORE_WHITE in possible:
            return Goban.SCORE_WHITE
        else:
            return None
        
        

    def _make_default_draw_codes(self):
        ret = []

        for pos in range(len(self.board)):
            if pos == 0:
                ret.append('ul')
            elif pos == self.board_size - 1:
                ret.append('ur')
            elif pos == self.board_size * self.board_size - self.board_size:
                ret.append('dl')
            elif pos == self.board_size * self.board_size - 1:
                ret.append('dr')
            elif pos in [60, 66, 72, 174, 180, 186, 288, 294, 300]:
                ret.append('h')
            elif pos < self.board_size - 1:
                ret.append('u')
            elif pos % self.board_size == 0:
                ret.append('l')
            elif pos > (self.board_size * self.board_size - self.board_size - 1):
                ret.append('d')
            elif pos % self.board_size == 18:
                ret.append('r')
            else:
                ret.append('m')

        return ret
            

    def _real_pos(self, pos):
        x,y = pos
        return x * self.board_size + y


    def _on_board(self, pos):
        return pos >= 0 and pos < self.board_size * self.board_size


    def _other_color(self, color):
        if color == Goban.BLACK:
            return Goban.WHITE
        elif color == Goban.WHITE:
            return Goban.BLACK


    def _neighbors(self, pos):
        neighbors = []
        if pos >= self.board_size:
            neighbors.append(pos - self.board_size)
        if pos <= self.board_size * self.board_size - self.board_size - 1:
            neighbors.append(pos + self.board_size)
        if pos % self.board_size != 0:
            neighbors.append(pos - 1)
        if (pos + 1) % self.board_size != 0:
            neighbors.append(pos + 1)

        return neighbors


    # Convert an sgf vector to a move tuple
    def _sgf_to_move(self, move):
        if move is None:
            return None

        x,y = move
        new_x = self.board_size - 1 - x
        return (new_x, y)


    # Convert a 1-dimensional position to an sgf move
    def _pos_to_sgf(self, pos):
        x = self.board_size - 1 + (pos / self.board_size)
        y = pos % self.board_size
        return (x,y)