import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))

from knowledge_graph import KG
from expert_policy import ExpertPolicy
from simulation.environment.unity_environment import UnityEnvironment

Task_set = [
    "Put apple on kitchencounter",
    "Put banana to coffeetable",
    "Rearrange book to bedroom",
    "Put mug on microwave",
    "Rearrange plate to kitchen",
    "Turn on tv",
    "Turn on radio",
    "Turn on microwave",
    "Turn on stove",
    "Turn on coffeemaker",
    "Turn off tv",
    "Turn off radio",
    "Turn off microwave",
    "Turn off stove",
    "Turn off coffeemaker",
    "Assemble apple and plate",
    "Assemble towel and barsoap",
    "Assemble towel and plate",
    "Assemble cutleryfork and cutleryknife",
    "Assemble cereal and milk",
    "Clean towel",
    "Clean apple",
    "Clean plate",
    "Clean cutleryfork",
    "Clean mug",
    "Place towel in washingmachine",
    "Place apple in fridge",
    "Place book in bookshelf",
    "Place notes in closet",
    "Place fryingpan in stove",
]

if __name__ == "__main__":
    env = UnityEnvironment(url='115.145.175.198')
    obs = env.reset(environment_id=0)
    kg = KG(env.get_position_graph())
    kg.add(obs['visible_graph'], 0, use_refinement=True)
    kg.add(obs['agent_graph'], 0, use_refinement=True)
    policy = ExpertPolicy(Task_set[0])

    print(policy.predict(kg))

