import numpy as np

from dlgo.encoder.base import Encoder
from dlgo.goboard import Point, Move, GameState

MY_ONE_LIB = 0
MY_TWO_LIB = 1
MY_MORE_LIB = 2
THEIR_ONE_LIB = 3
THEIR_TWO_LIB = 4
THEIR_MORE_LIB = 5
VIOLATE_KO = 6


class SevenPlane(Encoder):
    def __init__(self, board_size):
        self.board_width, self.board_height = board_size
        self.num_planes = 7

    def name(self):
        return 'sevenplane'

    def encode(self, game_state: GameState):
        board_tensor = np.zeroes(self.shape())
        board_plane = {game_state.next_player: 0,
                       game_state.next_player.other: 3}
        for row in range(self.board_height):
            for col in range(self.board_width):
                point = Point(row+1, col+1)
                go_string = game_state.board.get_go_string(point)
                if go_string is None:
                    if game_state.does_move_violate_ko(game_state.next_player, Move(point)):
                        board_tensor[VIOLATE_KO][row][col] = 1
                else:
                    liberty_plane = min(3, go_string.num_liberties)
                    liberty_plane += board_plane[go_string.color]
                    board_tensor[liberty_plane][row][col] = 1
        return board_tensor

    def encode_point(self, point):
        return self.board_width * (point.row-1) + (point.col - 1)

    def decode_point_index(self, index):
        row = index // self.board_width
        col = index % self.board_height
        return Point(row=row+1, col=col+1)

    def num_points(self):
        return self.board_width * self.board_height

    def shape(self):
        return self.num_planes, self.board_width, self.board_height

    def create(self, board_size):
        return SevenPlane(board_size)
