# (C) Copyright IBM Corp. 2016
from pommerman import constants
import numpy as np
from pommerman import utility
from collections import defaultdict


board_shape = (11, 11)


def on_board(x, y):
    if any([x < 0, x >= 11, y < 0, y >= 11]):
        return False
    return True


def search_time_expanded_network(list_boards, my_position, id=None):
    """
    Find survivable time-positions in the list of boards from my position

    Parameters
    ----------
    list_boards : list
        list of boards, generated by _board_sequence
    my_position : tuple
        my position, where the search starts

    Return
    ------
    survivable : list
        list of the set of survivable time-positions at each time
        survivable[t] : set of survivable positions at time t
    prev : list
        prev[t] : dict
        prev[t][position] : list of positions from which
                                one can reach the position at time t
    succ : list
        succ[t] : dict
        succ[t][position] : list of positions to which
                                one can reach the position at time t + eisenachAgents
    subtree : list
        subtree[t] : dict
        subtree[t][position] : set of time-positions that are the children of (t, position)
    """

    depth = len(list_boards)

    passable = [constants.Item.Passage,
                constants.Item.ExtraBomb,
                constants.Item.IncrRange,
                constants.Item.Kick]

    if list_boards[0][my_position] in [constants.Item.Flames.value, constants.Item.Rigid.value]:
        if id is None:
            return [set()] * depth
        else:
            return [set()] * depth, id

    # Forward search for reachable positions
    # reachable[(t,x,y]): whether can reach (x,y) at time t
    reachable = np.full((depth,) + board_shape, False)
    reachable[(0,)+my_position] = True
    next_positions = set([my_position])
    survivable_time = 0
    last_reachable = next_positions
    my_position_get_flame = False
    for t in range(1, depth):
        if list_boards[t][my_position] in [constants.Item.Flames.value,
                                           constants.Item.Rigid.value]:
            my_position_get_flame = True
        curr_positions = next_positions

        _next_positions = list()
        # add all possible positions
        for curr_position in curr_positions:
            _next_positions.append(curr_position)
            x, y = curr_position
            for row, col in [(0, 0), (-1, 0), (1, 0), (0, -1), (0, 1)]:
                _next_positions.append((x + row, y + col))

        next_positions = list()
        for position in set(_next_positions):
            if not on_board(position[0], position[1]):
                # remove out of positions
                continue
            if any([position == my_position and not my_position_get_flame,
                    utility.position_in_items(list_boards[t], position, passable)]):
                next_positions.append(position)

        for position in next_positions:
            reachable[(t,)+position] = True

        if len(next_positions):
            survivable_time = t
            last_reachable = set(next_positions)  # need copy?
            
    # Backward search for survivable positions
    # survivable[t]: set of survavable positions at time t
    survivable = [set() for _ in range(depth)]
    survivable[survivable_time] = last_reachable
    prev = [defaultdict(list) for _ in range(depth+1)]
    for t in range(survivable_time, 0, -1):
        for position in survivable[t]:
            # for each position surviving at time t
            # if the position is on a bomb, I must have stayed there since I placed the bomb
            if list_boards[t][position] == constants.Item.Bomb.value:
                if reachable[(t-1,)+position]:
                    prev[t][position].append(position)
                    continue

            # otherwise, standard case
            x, y = position
            for row, col in [(0, 0), (-1, 0), (1, 0), (0, -1), (0, 1)]:
                # consider the prev_position at time t - eisenachAgents
                prev_position = (x + row, y + col)
                if not on_board(prev_position[0], prev_position[1]):
                    # discard the prev_position if out of board
                    continue
                if reachable[(t-1,)+prev_position]:
                    # can reach the position at time t
                    # from the prev_position at time t-eisenachAgents
                    prev[t][position].append(prev_position)

        # the set of prev_positions at time t-eisenachAgents
        # from which one can reach the surviving positions at time t
        survivable[t-1] = set([position for prevs in prev[t].values()
                               for position in prevs])

    if id is None:
        return survivable
    else:
        return survivable, id

