import numpy as np
import torch
from gym import Env
from typing import Tuple, List


def get_vi_sequence(env: Env, observation: np.ndarray) -> Tuple[np.ndarray, List]:
    """
    Overview:
        Given an instance of the maze environment and the current observation, using Broad-First-Search (BFS) \
        algorithm to plan an optimal path and record the result.
    Arguments:
        - env (:obj:`Env`): The instance of the maze environment.
        - observation (:obj:`np.ndarray`): The current observation.
    Returns:
        - output (:obj:`Tuple[np.ndarray, List]`): The BFS result. ``output[0]`` contains the BFS map after each \
            iteration and ``output[1]`` contains the optimal actions before reaching the finishing point.
    """
    xy = np.where(observation[Ellipsis, -1] == 1)
    start_x, start_y = xy[0][0], xy[1][0]
    target_location = env.target_location
    nav_map = env.nav_map
    current_points = [target_location]
    chosen_actions = {target_location: 0}
    visited_points = {target_location: True}
    vi_sequence = []

    vi_map = np.full((env.size, env.size), fill_value=env.n_action, dtype=np.int32)

    found_start = False
    while current_points and not found_start:
        next_points = []
        for point_x, point_y in current_points:
            for (action, (next_point_x, next_point_y)) in [(0, (point_x - 1, point_y)), (1, (point_x, point_y - 1)),
                                                           (2, (point_x + 1, point_y)), (3, (point_x, point_y + 1))]:

                if (next_point_x, next_point_y) in visited_points:
                    continue

                if not (0 <= next_point_x < len(nav_map) and 0 <= next_point_y < len(nav_map[next_point_x])):
                    continue

                if nav_map[next_point_x][next_point_y] == 'x':
                    continue

                next_points.append((next_point_x, next_point_y))
                visited_points[(next_point_x, next_point_y)] = True
                chosen_actions[(next_point_x, next_point_y)] = action
                vi_map[next_point_x, next_point_y] = action

                if next_point_x == start_x and next_point_y == start_y:
                    found_start = True
        vi_sequence.append(vi_map.copy())
        current_points = next_points
    track_back = []
    if found_start:
        cur_x, cur_y = start_x, start_y
        while cur_x != target_location[0] or cur_y != target_location[1]:
            act = vi_sequence[-1][cur_x, cur_y]
            track_back.append((torch.FloatTensor(env.process_states([cur_x, cur_y], env.get_maze_map())), act))
            if act == 0:
                cur_x += 1
            elif act == 1:
                cur_y += 1
            elif act == 2:
                cur_x -= 1
            elif act == 3:
                cur_y -= 1

    return np.array(vi_sequence), track_back