import json

from data_gen_dependencies.omni_robot import IsaacSimRpcRobot
from data_gen_dependencies.omniagent import Agent

def generate_data(task_file, client_host, use_recording=True):
    task = json.load(open(task_file))
    robot_cfg = task["robot"]["robot_cfg"]
    stand = task["robot"]["stand"]
    robot_position = task["robot"]["init_position"]
    robot_rotation = task["robot"]["init_rotation"]
    scene_usd = task["scene_usd"]

    robot = IsaacSimRpcRobot(robot_cfg=robot_cfg, scene_usd=scene_usd,
                             client_host=client_host,
                             position=robot_position, rotation=robot_rotation, stand_type=stand["stand_type"],
                             stand_size_x=stand["stand_size_x"], stand_size_y=stand["stand_size_y"])
    agent = Agent(robot)
    render_semantic = False
    if "render_semantic" in task["recording_setting"]:
        render_semantic = task["recording_setting"]["render_semantic"]
    agent.run(task_file=task_file,
              camera_list=task["recording_setting"]["camera_list"],
              use_recording=use_recording,
              workspaces=task['scene']['function_space_objects'],
              fps=task["recording_setting"]["fps"],
              render_semantic=render_semantic,
              )
    print("job done")
    robot.client.Exit()