import json
# Python imports.
import random
from collections import defaultdict

# Other imports
from simple_rl.agents.AgentClass import Agent

class RandomVHAgent(Agent):

    def __init__(self, actions, actions_available, get_action_space, can_perform_action, get_graph, name=""):
        name = "Random" if name == "" else name
        self.actions_available = actions_available
        self.get_action_space = get_action_space
        self.can_perform_action = can_perform_action
        self.restriction_dict = None
        # self.get_graph = get_graph
        
        Agent.__init__(self, name=name, actions=actions)
    
    def populate_restriction_dict(self, state):
        self.restriction_dict = defaultdict(set)
        for node in state.data[0]["nodes"]:
            if 'CONTAINERS' in node["properties"]:
                self.restriction_dict["CONTAINERS"].add(node["class_name"])
            if 'SURFACES' in node["properties"]:
                self.restriction_dict["SURFACES"].add(node["class_name"])
            if 'GRABBABLE' in node["properties"]:
                self.restriction_dict["GRABBABLE"].add(node["class_name"])
            if "CAN_OPEN" in node["properties"]:
                self.restriction_dict["CAN_OPEN"].add(node["class_name"])


    def act(self, state, reward):
        # with open('s2.json', 'w') as f:
        #     json.dump(state.data, f, indent=4)

        if self.restriction_dict is None:
            # This only needs to be done once.
            self.populate_restriction_dict(state=state)
            # print(self.restriction_dict)
            restriction_dict_list = {k: list(v) for k, v in self.restriction_dict.items()}
            with open('restrict_dict.json', 'w') as f:
                json.dump(restriction_dict_list, f, indent=4)

        actionable_object_ids = self.get_action_space()[0]   # Importantly, action_space_ids is a list of *nodes* of objects that the agent can presently interact with!
        # It also include the list of rooms, which can be walked to.
        # This variable is actually a list of object ids, which should correspond to ids in the actual state.

        state_graph = state.data[0]
        # state_graph is a dictionary with keys "edges" and "nodes". The value for "nodes" is a list of dictionaries, each corresponding to a single node.
        actionable_objects_id_to_node = {node_dict['id']: node_dict for node_dict in state_graph['nodes'] if node_dict['id'] in actionable_object_ids}
        close_objects = [edge['to_id'] for edge in state_graph['edges'] if edge['from_id'] == 1 and edge['relation_type'] == 'CLOSE']

        # Our goal is to select an action that we can apply to a node whose ID is in action_space_ids.
        action_str = None
        while action_str is None:
            action_name = random.choice(self.actions_available)
            object_ids = []
            if action_name == "grab":
                # Visible objects whose classes are listed in grabbable
                object_ids = [id for id, object in actionable_objects_id_to_node.items() if object['class_name'] in self.restriction_dict["GRABBABLE"] and id in close_objects]
            elif action_name in ('walk', 'run'):
                object_ids = list(actionable_objects_id_to_node.keys())
            elif action_name == "open" or action_name == "close":
                object_ids = [id for id, object in actionable_objects_id_to_node.items() if object['class_name'] in self.restriction_dict["CAN_OPEN"] and id in close_objects]
            elif action_name.startswith('put'): # TODO: Probably need to re-write this
                action_name = random.choice(["putin", "putback"])
                obj_type = "CONTAINERS" if action_name == "putin" else "SURFACES"
                object_ids = [id for id, object in actionable_objects_id_to_node.items() if object['class_name'] in self.restriction_dict[obj_type] and id in close_objects]
            if len(object_ids) == 0:
                continue
            # The actions documentation is available here: http://virtual-home.org/documentation/master/kb/actions.html
            # This code currently skips a bunch of actions. See *SimpleVHEnv.py around line 45 and utils.py line 30

            object_id = random.choice(object_ids)
            action_str = self.can_perform_action(action_name, object_id, 1, state_graph, teleport=False)
        # print(action_str)
        return action_str
    
    # We may need a reset method to empty the restriction_dict
