import gym
import diffgro.environments.metaworld.envs.register
import diffgro.environments.metaworld_complex.envs.register


def make_env(domain_name, task_name, seed=777):
    if domain_name == "metaworld":
        env = gym.make(task_name, seed=seed)
    elif domain_name == "metaworld_complex":
        task_list = task_name.split("-")[:-1]

        # Actual env list will be given as argument, instead of task_name (therefore placeholder is used)
        if task_list[-1] == "variant":
            env = gym.make("complex-variant-v2", seed=seed, task_list=task_list[:-1])
        else:
            env = gym.make("complex-v2", seed=seed, task_list=task_list)

    return env


# alias for task to object
def task_to_object(domain_name, task_name):
    if domain_name == 'metaworld' and 'drawer' in task_name:
        return 'handle of drawer', "corner", \
            "In the image, the red robot arm is attempting to close the green drawer with white drawer handle. The robot's arm has a cylindrical blue and white tool attenched to its end, which is called as an end effector. The background is an indoor setting with a wooden floor and a gray wall.", \
            "is the robot arm is contacting to the handle of drawer?"
    elif domain_name == 'metaworld' and 'window' in task_name:
        return 'handle of window', "corner", \
            "In the image, the red robot arm with multiple joints is attempting to close the window with a small frame and two glass panes. The handle of the window is designed to slide the window closed. The robot's arm has a cylindrical blue and white tool attenched to its end, which is called as an end effector. The background is an indoor setting with a wooden floor and a gray wall.", \
            "is the robot arm's end effector is holding the handle of the window?" # with its side to the handle of window?"
    elif domain_name == 'metaworld' and 'faucet' in task_name:
        return 'handle of faucet', "corner3", \
            "In the image, the red robot arm is attempting to open the faucet. The faucet is visible with a distinct red faucet handle. The robot's arm has a cylindrical blue and white tool attenched to its end, which is called as an end effector. The background is an indoor setting with a wooden floor and a gray wall.", \
            "is the robot arm's end effector is contacting with its side to the handle of faucet?"
    elif 'metaworld_complex' == domain_name: # close_drawer_and_push_puck_and_peg_insert_and_press_button
        return "handle of drawer", "corner", \
            "In the image, the red robot arm is attempting to close the green drawer with white drawer handle. The robot's arm has a cylindrical blue and white tool attenched to its end, which is called as an end effector. The background is an indoor setting with a wooden floor and a gray wall.", \
            "is the robot arm is contacting to the handle of drawer?"
    else:
        raise NotImplementedError

