import copy
import numpy as np
import rlang
from rlang.grounding import ActionReference, MDPObject, MDPObjectGrounding, Feature, ConstantGrounding, Domain, Predicate, MDPClassGrounding, Proposition, Plan
from rlang import DictState
from grounding import get_stable_knowledge
import matplotlib.pyplot as plt
from functools import lru_cache
from simple_rl.tasks.gym.GymStateClass import GymState

class VHStateFeaturizer:    # TODO: Rewrite this for VH
    def __init__(self, knowledge):
        self.knowledge = knowledge
        self.skill_names_reversed = None


    def generate_rlang_objects(self, state):
        state = state[0]
        
        nodes = state['nodes']
        edges = state['edges']

        objs = {}

        # print(nodes)

        for n in nodes:
            if n['class_name'].capitalize() not in self.knowledge:
                print(f"Don't have class for object {n['class_name']}, it was not listed in get_stable_knowledge")
                continue
            cls = self.knowledge[n['class_name'].capitalize()]
            obj = cls(name=f"{n['class_name']}_{n['id']}", id=n['id'])
            objs[f"{n['class_name']}_{n['id']}"] = MDPObjectGrounding(obj=obj, domain=Domain.STATE)

        return objs

    
    def get_state_from_complete_predictions(self, state, predictions):
        new_state = copy.deepcopy(state[0])
        
        nodes = state['nodes']
        edges = state['edges']

        # inventory = list(inventory)
        # agent_info = list(agent_info)

        for k, v in predictions.items():
            if v == {}:
                continue
            # v = int(list(v.keys())[0])  # This is likely not an int anymore, it's likely a predicate's value

            print(k, v)

            return

            if OBJECT_TO_IDX[k.grounding.obj.name] == 10: # if agent
                if k.attribute_chain == ['x']:
                    agent_info[0] = v
                elif k.attribute_chain == ['y']:
                    agent_info[1] = v
                elif k.attribute_chain == ['dir']:
                    agent_info[2] = v
            elif OBJECT_TO_IDX[k.grounding.obj.name] == 4: # if door
                print(k.grounding.obj)
                if k.attribute_chain == ['is_open']:
                    raise NotImplementedError
                # TODO: Modify door state in the image
            
            # TODO: Extend this for objects other than agent.



def get_primitives_for(state_tuple, env=None):
    knowledge = get_stable_knowledge()
    statefeaturizer = VHStateFeaturizer(knowledge=knowledge)
    knowledge.update(statefeaturizer.generate_rlang_objects(state_tuple))
    # print(list(knowledge.keys()))
    # print("Mission:", env.env.mission)
    return list(knowledge.keys())


def get_knowledge_from_file(state_obj, filename="rlang_advice/pie_fridge.rlang", env=None):
    # state_dict = state_obj.data
    state_dict = state_obj
    knowledge = get_stable_knowledge()
    statefeaturizer = VHStateFeaturizer(knowledge=knowledge)
    knowledge.update(statefeaturizer.generate_rlang_objects(state_dict))
    knowledge = rlang.parse_file(filename, knowledge)
    print(list(knowledge.keys()))
    return knowledge, statefeaturizer
