import cv2
from virtualhome.simulation.environment.unity_environment import UnityEnvironment

if __name__ == "__main__":
    env = UnityEnvironment(base_port=8080, observation_types=['image'], num_agents=1)
    # env.relaunch()
    env.reset(environment_id=0)
    # entire_graph = env.get_graph()
    # for node in entire_graph['nodes']:
    #     if node['category'] not in ['Rooms', 'Walls', 'Ceiling', 'Lamps', 'Characters', 'Furniture', 'Floor', 'Decor', 'Floors', 'Doors', 'Windows']:
    #         print(node['category'], node['class_name'], node['prefab_name'], node['properties'])
    #         print(node)
    #
    # # print(env.get_graph())
    # print(env.get_action_space())
    # print(env.actions_available)

    while True:
        agent_action = input()
        env.step(agent_action)
        img = env.get_observations(info={'mode':'normal'})[0]
        print(env.get_visible_graph(mode='simple'))
        print(env.get_visible_objects())

        cv2.imshow('!!!', img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()