import random

import matplotlib.pyplot as plt
import networkx as nx
import torch.nn as nn

from src.gym_minigrid.minigrid import IDX_TO_COLOR, IDX_TO_OBJECT, STATE_TO_IDX

IDX_TO_STATE = dict(zip(STATE_TO_IDX.values(), STATE_TO_IDX.keys()))

ACTION_TO_TEXT = {
    0: "turns left",
    1: "turns right",
    2: "moves forward",
    3: "picks up",
    4: "drops",
    5: "toggles",
    6: "is done"
}


def describe_observation(obs_image):
    """
    Returns a description of the observation (list of objects in the environment).
    For now it has what objects the agent is observing.
    TODO: Include description of relative location to agent position:
    FYI "agent_pos=(self.agent_view_size // 2, self.agent_view_size - 1)," => (3, 6)
    :returns
    a list of tuples representing all the objects the agent sees in its view and the color (obj_type, color)
    """
    objects = []
    for i in range(obs_image.shape[0]):
        for j in range(obs_image.shape[1]):
            type_idx, color_idx, state = obs_image[i, j]
            obj_type = IDX_TO_OBJECT[type_idx]
            color = IDX_TO_COLOR[color_idx]
            state = IDX_TO_STATE[
                state]  # STATE used only for some cases like the door, where an object have additional state descriptors (open, closed, locked)
            if obj_type == 'door':
                obj_type = state + ' ' + obj_type
            if obj_type == 'unseen':
                continue
            objects.append((obj_type, color))
    return objects


def describe_sequences(observations, observations_sequence):
    """
    A function used to provide linguistic descriptions of sequences
    :param observations: A dictionary of observations. For each observation, you can get their feature representations by calling the method all_obs[obs_id].describe_observation()
    :param observations_sequence: An array of tuples each tuple has an obs_id and the action that follows from this observation.
    :return:
    """
    sequences_as_text = ""
    for observation, action in observations_sequence:
        sentences, description = observations[observation].describe_lang_observatoion()
        if action is None:
            sequences_as_text += description + "Then, the agent does nothing. \n \n"
            continue
        else:
            sequences_as_text += description + "Then, the agent " + ACTION_TO_TEXT[action]

        if action in [3, 4, 5] and (0, 1, 0) in sentences:
            sequences_as_text += " the " + sentences[(0, 1, 0)]

        sequences_as_text += ". \n \n"
    return sequences_as_text


class DQN(nn.Module):
    def __init__(self, input_shape, num_actions):
        super(DQN, self).__init__()
        self.input_shape = input_shape
        self.num_actions = num_actions
        self.layers = nn.Sequential(
            nn.Linear(input_shape[0], 64), nn.ReLU(), nn.Linear(64, self.num_actions)
        )

    def forward(self, x):
        return self.layers(x)


def hash_img(img):
    """
    Convert images to a 1-D string to be able to use them in dict as keys.
    Args:
        img: numpy ndarray of integers
    """
    return "".join([str(i) for i in list(img.reshape(-1))])


class ObsNode:
    def __init__(self, obs_id, agent_pos=None, agent_dir=None):
        """
        An observation node to be used in a linked list to represent the observations
        in an agent's path and their outgoing actions and subsequent observation.
        Args:
            obs_id: observation id, a string that starts with the letter O followed by an index corresponding
            to the order of its initiation.
        """
        self.obs_id = obs_id
        self.action = None
        self.q_value = 0
        self.next = None
        self.agent_dir = agent_pos
        self.agent_pos = agent_dir


class AgentPath:
    def __init__(self, name, p_type, description):
        """
        An agent's path contains a chain of observations linked by actions.
        It can be the main agent path or a path generated by
        Args:
            name: A unique path name to serve as a path identifier
            p_type: A type to differentiate main agents path from generated explanation paths, values are either main path or explanation path.
            description: A description for the path containing extra information about it.
        """
        self.name = name  # A unique name of the path to differentiate from others.
        # For example, main_path, exp1, exp2.
        self.path_type = p_type  # either [explanation or main].
        self.description = description  # A description of the path and what it represents.
        if p_type != 'main':
            self.edge_color = '#%06X' % random.randint(0, 0xFFFFFF)
        else:
            self.edge_color = '#000000'
        self.head = None
        self.final_reward = 0

    def init_node(self, obs1, agent_pos=None, agent_dir=None):
        """
        Initialize node, this method is typically used for the first observation in a path,
        were the subsequent observation and action are unkown yet.
        Args:
            obs1: The observation id, a string generated in the ExplanationGraph object.
            :param obs1:
            :param agent_pos:
            :param agent_dir:
        """
        n = ObsNode(obs1, agent_pos, agent_dir)
        if not self.head:
            self.head = n
        return n

    def add_edge(self, obs1_node, obs2_id, new_action, new_q, agent_pos=None, agent_dir=None):
        """
        Args:
            obs1_node:
            obs2_id:
            new_action:
            :param obs1_node:
            :param obs2_id:
            :param new_action:
            :param new_q:
            :param agent_pos:
            :param agent_dir:
        """
        obs1_node.action = new_action
        obs1_node.q_value = new_q
        obs2_node = ObsNode(obs2_id, agent_pos, agent_dir)

        obs1_node.next = obs2_node
        return obs2_node

    def get_action(self, obs1_id, obs2_id):
        """
        Args:
            obs1_id:
            obs2_id:
        """
        current = self.head
        while current is not None:
            if current.obs_id == obs1_id and current.next.obs_id == obs2_id:
                return current.action
            else:
                current = current.next
        return None

    def get_edges(self):
        edges = []
        current = self.head
        while current is not None:
            next_e = current.next
            if next_e is not None:
                pair = (current.obs_id, next_e.obs_id)
                edges.append(pair)
            current = next_e
        return edges

    def get_edges_with_q_values(self):
        edges = []
        current = self.head
        while current is not None:
            next_e = current.next
            if next_e is not None:
                pair = (
                    current.obs_id, next_e.obs_id, current.action, current.q_value, current.agent_pos, current.agent_dir)
                edges.append(pair)
            current = next_e
        return edges

    def get_q_values(self):
        qs = []
        current = self.head
        while current is not None:
            qs.append((current.obs_id, current.q_value))
            current = current.next
        return qs

    def printpath(self):
        current = self.head
        while current is not None:
            print(current.obs_id)
            #  TODO What else to print??
            current = current.next

    def get_path_array(self):
        paths_arr = []
        current = self.head
        while current is not None:
            try:
                action = current.action.action
            except AttributeError:
                action = None
            paths_arr.append((current.obs_id, action, current.q_value, current.agent_pos, current.agent_dir))
            current = current.next
        return paths_arr


class Observation:
    def __init__(self, unique_id, obs, alias):
        """
        Args:
            unique_id:
            obs:
        """
        self.obs_id = unique_id  # Given an observation image generate a unique ID (Unique per graph).
        self.alias = alias
        self.obs_image = obs
        self.agent_position = None
        self.description = None

    def describe_lang_observatoion(self):
        # Return a textual description of the observation
        # FYI "agent_pos=(self.agent_view_size // 2, self.agent_view_size - 1)," => (3, 6)
        # [-] [-] [-] [-] [-] [-] [-]
        # [-] [-] [-] [-] [-] [-] [-]
        # [-] [-] [-] [-] [-] [-] [-]
        # [-] [-] [-] [-] [-] [-] [-]
        # [-] [12] [11] [10] [9] [8] [-]
        # [-] [13] [4] [3] [2] [7] [-]
        # [-] [14] [5] [*] [1] [6] [-]
        # East of the agent, North east of the agent, North of the agent, North west of the agent, west of the agent
        # N spot to the east ...,   N spot to the north and N spots to the east ...
        sentences = dict()  # Each key is a tuple (W,N,E)
        agent_pos = (self.obs_image.shape[0] // 2, self.obs_image.shape[1] - 1)
        for i in range(self.obs_image.shape[0]):
            for j in range(self.obs_image.shape[1]):
                type_idx, color_idx, state = self.obs_image[i, j]
                obj_type = IDX_TO_OBJECT[type_idx]
                color = IDX_TO_COLOR[color_idx]
                state = IDX_TO_STATE[state]
                if obj_type == 'door':
                    obj_type = state + ' ' + obj_type
                if obj_type == 'unseen':
                    continue
                    # West: i < agent_i
                    # East: i > agent_i
                    # North: j < agent_j
                if obj_type == 'empty':
                    sentences[
                        (max(0, agent_pos[0] - i), agent_pos[1] - j, max(i - agent_pos[0], 0))] = obj_type + " space"
                else:
                    sentences[
                        (max(0, agent_pos[0] - i), agent_pos[1] - j, max(i - agent_pos[0], 0))] = color + " " + obj_type

        proper_description = ""
        if (0, 0, 0) in sentences:
            if 'key' in sentences[(0, 0, 0)]:
                proper_description += "The agent has the " + sentences[(0, 0, 0)] + ". "
            else:
                proper_description += "The agent is at the " + sentences[(0, 0, 0)] + ". "
        # Count sum of each tuple manhattan distance.
        distance_counter = dict([])
        for k in sentences:
            s_k = sum(list(k))
            if s_k > 0:
                if s_k in distance_counter:
                    distance_counter[sum(list(k))].append(k)
                else:
                    distance_counter[sum(list(k))] = [k]

        slots = ['slots', 'spaces']
        slot = ['slot', 'space']

        for k in sorted(distance_counter.keys()):
            for t in distance_counter[k]:
                west = t[0]
                north = t[1]
                east = t[2]
                if sentences[t][0] in ["a", "e", "i", "o", "u"]:
                    proper_description += "There is an " + sentences[t] + " located "
                else:
                    proper_description += "There is a " + sentences[t] + " located "
                if north == 1:
                    if east == 0 and west == 0:
                        proper_description += str(north) + " " + random.choice(slot) + " in front of the agent. "
                    else:
                        proper_description += str(north) + " " + random.choice(slot) + " in front of the agent and "

                if north > 1:
                    if east == 0 and west == 0:
                        proper_description += str(north) + " " + random.choice(slots) + " in front of the agent. "
                    else:
                        proper_description += str(north) + " " + random.choice(slots) + " in front of the agent and "
                if east == 1:
                    proper_description += str(east) + " " + random.choice(slot) + " to the agent's right. "
                if west == 1:
                    proper_description += str(west) + " " + random.choice(slot) + " to the agent's left. "
                if east > 1:
                    proper_description += str(east) + " " + random.choice(slots) + " to the agent's right. "
                if west > 1:
                    proper_description += str(west) + " " + random.choice(slots) + " to the agent's left. "
                # proper_description += "\n"
        return sentences, proper_description

    def describe_observation(self):
        """
        Returns a description of the observation (list of objects in the environment).
        For now it has what objects the agent is observing.
        TODO: Include description of relative location to agent position:
        FYI "agent_pos=(self.agent_view_size // 2, self.agent_view_size - 1)," => (3, 6)
        :returns
        a list of tuples representing all the objects the agent sees in its view and the color (obj_type, color)
        """
        return describe_observation(self.obs_image)

    def visualize(self, obs_env, tile_size=16):
        """
        Visualize the observation for humans (instead of translating to language, translate to image)
        Args:
            obs_env:
            :param obs_env:
            :param tile_size:
        """
        obs = obs_env.get_obs_render(self.obs_image, tile_size)
        # _, ax = plt.subplots(nrows=1, ncols=1, figsize=(6, 6),
        #                     subplot_kw={'xticks': [], 'yticks': []})
        # a = ax.imshow(obs, interpolation='bilinear')
        # a.set_data(obs)
        # plt.tight_layout()
        # plt.show()
        return obs


class Action:
    def __init__(self, action_id, action, reward, done):
        """
        Args:
            action_id:
            action:
            reward:
            done:
        """
        self.action_id = action_id  # Autogenerated action ID, unique per graph.
        self.action = action
        self.reward = reward
        self.done = done
        self.probability = None  # TODO: Get this from agent.

    @property
    def action(self):
        return self.__action

    @action.setter
    def action(self, a):
        # TODO: Check if action in list of valid actions?
        """
        Args:
            a:
        """
        self.__action = a


class ExplanationGraph:

    def __init__(self, env=None):
        self.graph = nx.MultiDiGraph()
        self.image_to_alias = dict()
        self.__observations = dict()
        self.action_ids = 0
        self.obs_id = 0
        self.paths = dict()
        self.env = env

    @property
    def observations(self):
        return self.__observations

    # NOT used
    @observations.setter
    def observations(self, episodes):
        # Given a list of observations add them to the list and create their unique ids.
        """
        Args:
            episodes:
        """
        pre_obs_id = None
        for episode in episodes:
            action = episode[0]
            obs = episode[1]
            reward = episode[2]
            done = episode[3]
            path_name = episode[5]
            pre_obs_id = self.add_observation(action, obs, reward, done, pre_obs_id, path_name, None, None)

    # class Observations(dict):
    #     def __init__(self, image_to_id):
    #         super().__init__()
    #         self.image_to_id = image_to_id
    #         self.__observations = dict()
    #
    #     def __getitem__(self, arg):
    #         item = self.image_to_id[arg]
    #         return super().__getitem__(item)

    def create_path(self, path_name, path_type, path_description=None):
        # TODO: Add checks to ensure correct path_type and unique path name.
        """
        Args:
            path_name:
            path_type:
            path_description:
        """
        self.paths[path_name] = AgentPath(path_name, path_type, path_description)
        return True

    def add_observation(self, action, obs, reward, done, pre_obs, path_name, q_value, agent_pos=None, agent_dir=None):
        # Add observation to observations list and make sure that the list has unique Ids.
        """
        Args:
            action:
            obs:
            reward:
            done:
            pre_obs:
            path_name:
            q_value:
            agent_pos:
            :param action:
            :param obs:
            :param reward:
            :param done:
            :param path_name:
            :param pre_obs:
            :param q_value:
            :param agent_pos:
            :param agent_dir:
        """
        hashed_img = hash_img(obs['image'])
        alias = self.get_img_id(hashed_img)
        unique_id = hashed_img
        if not alias:
            unique_id = hashed_img
            alias = "O" + str(self.obs_id)
            # generate id and put image in img_to_id dict and observations list
            observation = Observation(unique_id, obs['image'], alias)
            self.image_to_alias[hashed_img] = alias
            self.__observations[unique_id] = observation
            self.obs_id += 1

        # Update graph:
        if pre_obs:
            action_in = Action(self.action_ids, action, reward, done)
            self.action_ids += 1
            new_obs = self.paths[path_name].add_edge(pre_obs, unique_id, action_in, q_value, agent_pos, agent_dir)
        else:
            new_obs = self.paths[path_name].init_node(unique_id, agent_pos, agent_dir)
            new_obs.q_value = q_value

        return new_obs

    def construct_graph(self, paths_to_show=None):
        """
        Args:
            paths_to_show:
        """
        if paths_to_show is None:
            paths_to_show = ['main']
        self.graph = nx.MultiDiGraph()
        for path in paths_to_show:
            edges = self.paths[path].get_edges()  # a list of tuples.
            converted_edges = [(self.image_to_alias[e[0]], self.image_to_alias[e[1]]) for e in edges]
            self.graph.add_edges_from(converted_edges,
                                      color=self.paths[path].edge_color)

    def reset_graph(self):
        # Function to reset the graph and only include the main path.
        self.construct_graph(['main'])

    def draw_graph(self):
        # TODO Draw graph showing the simulation and the counterfactuals as different node types. 
        plt.subplot(111)

        colors = nx.get_edge_attributes(self.graph, 'color')
        labels = {k: k for k in self.graph.nodes}
        pos = nx.spring_layout(self.graph)
        nx.draw_networkx_nodes(self.graph, pos)
        nx.draw_networkx_labels(self.graph, pos, labels)

        ax = plt.gca()
        for e in self.graph.edges:
            ax.annotate("",
                        xy=pos[e[0]], xycoords='data',
                        xytext=pos[e[1]], textcoords='data',
                        arrowprops=dict(arrowstyle="<-", color=colors[e],
                                        shrinkA=5, shrinkB=5,
                                        patchA=None, patchB=None,
                                        connectionstyle="arc3,rad=rrr".replace('rrr', str(0.3 * e[2])
                                                                               ),
                                        ),
                        )
        plt.axis('off')
        plt.show()

    def get_img_id(self, hashed_img):
        """
        Args:
            hashed_img:
        """
        img_id = None
        if hashed_img in self.image_to_alias:
            # Retrieve obs with this id and don't add into list
            img_id = self.image_to_alias[hashed_img]
        return img_id


if __name__ == "__main__":
    # Main is in: observations_data_generation.py.
    pass
