import numpy as np
from itertools import permutations
from kaggle_environments.envs.hungry_geese.hungry_geese import Action, adjacent_positions
from args import *


def str_to_action(s):
    if s == 'NORTH':
        return Action.NORTH
    elif s == 'EAST':
        return Action.EAST
    elif s == 'SOUTH':
        return Action.SOUTH
    else:
        return Action.WEST


def get_board(obs, prev_actions, args):
    '''
    Channels of state:
        geese curr_heads  0,  1,  2,  3
        geese bodies      4,  5,  6,  7
        geese tips        8,  9, 10, 11
        geese prev_heads 12, 13, 14, 15
        food             16
    '''
    board = np.zeros((args.numAgents * 4 + 1, args.boardSize[0] * args.boardSize[1]), np.uint8)

    for i, goose in enumerate(obs.geese):
        # head position
        for head_pos in goose[:1]:
            board[0 + (i - obs.index) % args.numAgents, head_pos] = 1

        # tip position
        for tip_pos in goose[-1:]:
            board[args.numAgents + (i - obs.index) % args.numAgents, tip_pos] = 1

        # whole position
        for body_pos in goose[1:]:
            board[args.numAgents * 2 + (i - obs.index) % args.numAgents, body_pos] = 1

        # previous head position
        for head_pos in goose[:1]:
            if prev_actions[i] is not None:
                prev_action = str_to_action(prev_actions[i])
                opposite_action = Action.opposite(prev_action)
                prev_head_pos = adjacent_positions(head_pos, args.boardSize[1], args.boardSize[0])[opposite_action.value - 1]
                board[args.numAgents * 3 + (i - obs.index) % args.numAgents, prev_head_pos] = 1

    for food_pos in obs.food:
        board[-1, food_pos] = 1

    return board.reshape(-1, args.boardSize[0], args.boardSize[1])


def get_player_board(board, player, num_agents):
    new_board = board.copy()
    indices = np.arange(0, num_agents * 4, num_agents)
    tmp = new_board[indices]
    new_board[indices] = new_board[indices + player]
    new_board[indices + player] = tmp
    return new_board


def get_symmetries(board, player, num_agents, pi):
    boards = [(board, pi)]

    # 180 degree turn
    b = np.rot90(board, 2, (1, 2))
    p = [*pi[2:], *pi[:2]]
    boards.append((b, p))

    # flip ud
    b = np.flip(board, 1)
    p = pi[:]
    tmp = p[0]
    p[0] = p[2]
    p[2] = tmp
    boards.append((b, p))

    # flip lr
    b = np.flip(board, 2)
    p = pi[:]
    tmp = p[1]
    p[1] = p[3]
    p[3] = tmp
    boards.append((b, p))

    # change player indices
    indices = np.arange(0, num_agents * 4, num_agents)
    others = [i for i in range(num_agents) if i != player]
    for bb, p in boards[:4]:
        for pm in permutations(others, 3):
            b = bb.copy()
            for i, j in zip(others, pm):
                b[indices + i] = bb[indices + j]
            boards.append((b, p))

    return boards


def get_reward(obs, player, num_agents):
    alive = 0
    for goose in obs.geese:
        alive += 1 if len(goose) > 0 else 0

    if len(obs.geese[player]) > 0 and alive >= 2 and obs.step < 199:
        return None  # not finished

    rank = 1
    for i, goose in enumerate(obs.geese):
        if i == player:
            continue
        if len(goose) > len(obs.geese[player]):
            rank += 1
        elif len(goose) == len(obs.geese[player]):
            rank += 0.5
    return (num_agents + 1 - 2 * rank) / (num_agents - 1)


def get_valid_moves(state, player, action_size):
    moves = np.ones(action_size)

    if state[0].observation.step == 0:
        return moves

    prev_action = str_to_action(state[player]['action'])
    invalid_action = Action.opposite(prev_action)
    moves[invalid_action.value - 1] = 0
    return moves


def select_action(pi, prev_action):
    '''
    Args:
        pi: np.array(4)
        prev_action: one of 0, 1, 2, 3
    Return:
        action: str
    '''
    if prev_action is not None:
        prev_action = str_to_action(prev_action)
        invalid_action = Action.opposite(prev_action)
        pi[invalid_action.value - 1] = 0
        pi /= np.sum(pi)
    action_num = np.random.choice(len(pi), p=pi) + 1
    action = Action(action_num).name
    return action


def agent(obs, config, prev_actions=None):
    if prev_actions is None:
        try:
            prev_actions = agent.prev_actions
        except AttributeError:
            agent.prev_actions = [None] * args.numAgents
            prev_actions = agent.prev_actions

    player = obs.index

    board = get_board(obs, prev_actions, args)
    board = get_player_board(board, player, args.numAgents)

    pi, _ = agent.net.predict(board)

    # remove invalid action
    prev_action = prev_actions[player]
    if prev_action is not None:
        prev_action = str_to_action(prev_action)
        oppo_action = Action.opposite(prev_action).value
        pi[oppo_action - 1] = 0
        pi /= pi.sum()

    action = np.argmax(pi) + 1
    action = Action(action).name
    prev_actions[player] = action
    return action
