gtp_extensions.py

# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# extends gtp.py

import gtp
import sys
import sgf_wrapper
import itertools
import go
import coords

def parse_message(message):
    message = gtp.pre_engine(message).strip()
    first, rest = (message.split(" ", 1) + [None])[:2]
    if first.isdigit():
        message_id = int(first)
        if rest is not None:
            command, arguments = (rest.split(" ", 1) + [None])[:2]
        else:
            command, arguments = None, None
    else:
        message_id = None
        command, arguments = first, rest

    command = command.replace("-", "_") # for kgs extensions.
    return message_id, command, arguments


class KgsExtensionsMixin(gtp.Engine):

    def __init__(self, game_obj, name="gtp (python, kgs-chat extensions)", version="0.1"):
        super().__init__(game_obj=game_obj, name=name, version=version)
        self.known_commands += ["kgs-chat"]

    def send(self, message):
        message_id, command, arguments = parse_message(message)
        if command in self.known_commands:
            try:
                retval = getattr(self, "cmd_" + command)(arguments)
                response = gtp.format_success(message_id, retval)
                sys.stderr.flush()
                return response
            except ValueError as exception:
                return gtp.format_error(message_id, exception.args[0])
        else:
            return gtp.format_error(message_id, "unknown command: " + command)

    # Nice to implement this, as KGS sends it each move.
    def cmd_time_left(self, arguments):
        pass

    def cmd_showboard(self, arguments):
        return self._game.showboard()

    def cmd_kgs_chat(self, arguments):
        try:
            msg_type, sender, *text = arguments.split()
            text = " ".join(text)
        except ValueError:
            return "Unparseable message, args: %r" % arguments
        return self._game.chat(msg_type, sender, text)

class RegressionsMixin(gtp.Engine):
    def cmd_loadsgf(self, arguments):
        args = arguments.split()
        if len(args) == 2:
            file_, movenum = args
            movenum = int(movenum)
            print("movenum =", movenum, file=sys.stderr)
        else:
            file_ = args[0]
            movenum = None

        try:
            with open(file_, 'r') as f:
                contents = f.read()
        except:
            raise ValueError("Unreadable file: " + file_)

        try:
            # This is kinda bad, because replay_sgf is already calling
            # 'play move' on its internal position objects, but we really
            # want to advance the engine along with us rather than try to
            # push in some finished Position object.
            for idx, p in enumerate(sgf_wrapper.replay_sgf(contents)):
                print("playing #", idx, p.next_move, file=sys.stderr)
                self._game.play_move(p.next_move)
                if movenum and idx == movenum:
                    break
        except:
            raise

# Should this class blatantly reach into the game_obj and frob its tree?  Sure!
# What are private members?  Python lets you do *anything!*
class GoGuiMixin(gtp.Engine):
    """ GTP extensions of 'analysis commands' for gogui.
    We reach into the game_obj (an instance of the players in strategies.py),
    and extract stuff from its root nodes, etc.  These could be extracted into
    methods on the Player object, but its a little weird to do that on a Player,
    which doesn't really care about GTP commands, etc.  So instead, we just
    violate encapsulation a bit... Suggestions welcome :) """
    def __init__(self, game_obj, name="gtp (python, gogui extensions)", version="0.1"):
        super().__init__(game_obj=game_obj, name=name, version=version)
        self.known_commands += ["gogui-analyze_commands"]

    def cmd_gogui_analyze_commands(self, arguments):
        return "\n".join(["var/Most Read Variation/nextplay",
                          "var/Think a spell/spin",
                          "pspairs/Visit Heatmap/visit_heatmap",
                          "pspairs/Q Heatmap/q_heatmap"])

    def cmd_nextplay(self, arguments):
        return self._game.root.mvp_gg()

    def cmd_visit_heatmap(self, arguments):
        sort_order = list(range(self._game.size * self._game.size + 1))
        sort_order.sort(key=lambda i: self._game.root.child_N[i], reverse=True)
        return self.heatmap(sort_order, self._game.root, 'child_N')

    def cmd_q_heatmap(self, arguments):
        sort_order = list(range(self._game.size * self._game.size + 1))
        reverse = True if self._game.root.position.to_play is go.BLACK else False
        sort_order.sort(key=lambda i: self._game.root.child_Q[i], reverse=reverse)
        return self.heatmap(sort_order, self._game.root, 'child_Q')

    def heatmap(self, sort_order, node, prop):
        return "\n".join(["{!s:6} {}".format(
            coords.to_human_coord(coords.unflatten_coords(key)),
            node.__dict__.get(prop)[key])
            for key in sort_order if node.child_N[key] > 0][:20])

    def cmd_spin(self, arguments):
        for i in range(50):
            for j in range(100):
                self._game.tree_search()
            moves = self.cmd_nextplay(None).lower()
            moves = moves.split()
            colors = "bw" if self._game.root.position.to_play is go.BLACK else "wb"
            moves_cols = " ".join(['{} {}'.format(*z) 
                                   for z in zip(itertools.cycle(colors), moves)])
            print("gogui-gfx: TEXT", "{:.3f} after {}".format(self._game.root.Q, self._game.root.N), file=sys.stderr, flush=True)
            print("gogui-gfx: VAR", moves_cols, file=sys.stderr, flush=True)
        return self.cmd_nextplay(None)

class GTPDeluxe(KgsExtensionsMixin, RegressionsMixin, GoGuiMixin):
    pass