#!/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 row_strength(field, row):
    strength = 0
    bonus = 1
    for col in range(field.shape[1]):
        if field[row][col] > 0:
            bonus = bonus * 2
            strength += bonus
        else:
            bonus = 1
    return strength


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 has_top_neighbour(field, row, col):
    for r in range(0,row):
        if field[r][col] > 0:
            return True
    return False


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

    for r in range(field.shape[0]):
        for c in range(field.shape[1]):
            if has_top_neighbour(field, r, c):
                result -= 1000.0

    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")

    builder = flatbuffers.Builder()
    wire_protocol.CommandRequest.Start(builder)
    wire_protocol.CommandRequest.AddCommand(
            builder,
            wire_protocol.Command.Command.StartGame)
    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)

    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)
        if rsp.Error() != wire_protocol.CommandResponseError.CommandResponseError.OK:
           return

        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)
                if rsp.Error() != wire_protocol.CommandResponseError.CommandResponseError.OK:
                    return

            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)
                if rsp.Error() != wire_protocol.CommandResponseError.CommandResponseError.OK:
                    return


            if best_move.col_offset == 0:
                builder = flatbuffers.Builder()
                wire_protocol.CommandRequest.Start(builder)
                wire_protocol.CommandRequest.AddCommand(
                        builder, wire_protocol.Command.Command.Down)
                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 (rsp.Error() != wire_protocol.CommandResponseError.CommandResponseError.OK and
                    rsp.Error() != wire_protocol.CommandResponseError.CommandResponseError.INVALID_COMMAND):
                    return

if __name__ == "__main__":
    main()