competitive_tetris/bot_example.py

#!/usr/bin/python3

"""
Competitive tetris bot example
"""

import time
import random
import zmq
import flatbuffers
import numpy as np

import wire_protocol.Command
import wire_protocol.CommandRequest
import wire_protocol.CommandResponse
import wire_protocol.CommandResponseError

import game_state

def dummy_next_piece_func():
    return 0

def cell_strength(field, row, col):
    neighbour_count = 0
    for r in range(row-1, row+2):
        for c in range(col-1, col+2):
            if (r < 0 or r >= field.shape[0] or
                c < 0 or c >= field.shape[1]):
                # walls are neighbours
                neighbour_count += 1
            elif field[r][c] > 0:
                neighbour_count += 1
    if field[row][col]:
        return neighbour_count * row
    else:
        return -neighbour_count * row


def field_heuristic(field):
    result = 0.0
    for r in range(field.shape[0]):
        for c in range(field.shape[1]):
            result += cell_strength(field, r, c)
    return result


class ShiftAndRotate:
    def __init__(self, rotate_count, col_offset, value, state):
        self.rotate_count = rotate_count
        self.col_offset = col_offset
        self.value = value
        self.state = state


def shift_state(state, shift):
    if shift == 0:
        return state

    if shift < 0:
        for _ in range(abs(shift)):
            next_state = state.shift_left()
            if next_state is None:
                return None
            state = next_state
    else:
        for _ in range(shift):
            next_state = state.shift_right()
            if next_state is None:
                return None
            state = next_state

    return state


def create_shifts_and_rotates(state):
    shift_and_rotates = []

    for rotate_count in range(5):
        rotate_state = state
        for _ in range(rotate_count):
            rotate_state = rotate_state.rotate()
            if rotate_state is None:
                break

        if rotate_state is None:
            continue

        for col_offset in range(-game_state.FIELD_WIDTH, game_state.FIELD_WIDTH+1):

            next_state = shift_state(rotate_state, col_offset)
            if next_state is not None:
                final_state, is_running = next_state.land().advance()
                if is_running:
                    value = (field_heuristic(final_state.data.grid) +
                             final_state.data.score*10000.0)
                    shift_and_rotates.append(ShiftAndRotate(
                            rotate_count,
                            col_offset,
                            value,
                            final_state))
    return shift_and_rotates


def main():
    context = zmq.Context()
    socket = context.socket(zmq.REQ)
    socket.connect("tcp://localhost:5555")

    while True:
        builder = flatbuffers.Builder()
        wire_protocol.CommandRequest.Start(builder)
        wire_protocol.CommandRequest.AddCommand(
                builder,
                wire_protocol.Command.Command.GetGameState)
        req = wire_protocol.CommandRequest.End(builder)
        builder.Finish(req)
        socket.send(builder.Output())

        rsp_buf = socket.recv()
        rsp = wire_protocol.CommandResponse.CommandResponse.GetRootAs(rsp_buf)

        grid = rsp.GameState().GridAsNumpy()
        grid = grid.reshape((rsp.GameState().Rows(), rsp.GameState().Cols()))

        state = game_state.GameState()
        state.data = game_state.GameStateData(
                grid = grid,
                current_piece = rsp.GameState().CurrentPiece(),
                next_piece = rsp.GameState().NextPiece(),
                rotation = rsp.GameState().PieceRotation(),
                row = rsp.GameState().PieceRow(),
                col = rsp.GameState().PieceCol(),
                score = rsp.GameState().Score())
        state.next_piece_func = dummy_next_piece_func


        next_moves = create_shifts_and_rotates(state)
        if len(next_moves) > 0:
            best_moves = sorted(next_moves, key=lambda x: x.value, reverse=True)
            best_move = best_moves[0]

            command = wire_protocol.Command.Command.Rotate
            for _ in range(best_move.rotate_count):
                builder = flatbuffers.Builder()
                wire_protocol.CommandRequest.Start(builder)
                wire_protocol.CommandRequest.AddCommand(builder, command)
                req = wire_protocol.CommandRequest.End(builder)
                builder.Finish(req)
                socket.send(builder.Output())
                rsp_buf = socket.recv()
                rsp = wire_protocol.CommandResponse.CommandResponse.GetRootAs(rsp_buf)

            command = wire_protocol.Command.Command.ShiftRight
            if best_move.col_offset < 0:
                command = wire_protocol.Command.Command.ShiftLeft
            for _ in range(abs(best_move.col_offset)):
                builder = flatbuffers.Builder()
                wire_protocol.CommandRequest.Start(builder)
                wire_protocol.CommandRequest.AddCommand(builder, command)
                req = wire_protocol.CommandRequest.End(builder)
                builder.Finish(req)
                socket.send(builder.Output())
                rsp_buf = socket.recv()
                rsp = wire_protocol.CommandResponse.CommandResponse.GetRootAs(rsp_buf)

            builder = flatbuffers.Builder()
            wire_protocol.CommandRequest.Start(builder)
            wire_protocol.CommandRequest.AddCommand(
                    builder, wire_protocol.Command.Command.Land)
            req = wire_protocol.CommandRequest.End(builder)
            builder.Finish(req)
            socket.send(builder.Output())
            rsp_buf = socket.recv()
            rsp = wire_protocol.CommandResponse.CommandResponse.GetRootAs(rsp_buf)

if __name__ == "__main__":
    main()