from gym import envs

from diffgro.environments.metaworld.envs import *
from config.metaworld import *


########## Base Environments ##########

envs.register(
    id="button-press-v2",
    entry_point="diffgro.environments.metaworld.envs.env:MetaWorldEnv",
    kwargs={
        "env_name": "button-press-v2",
        "max_steps": 500,
    },
)

envs.register(
    id="door-open-v2",
    entry_point="diffgro.environments.metaworld.envs.env:MetaWorldEnv",
    kwargs={
        "env_name": "door-open-v2",
        "max_steps": 500,
    },
)


envs.register(
    id="drawer-close-v2",
    entry_point="diffgro.environments.metaworld.envs.env:MetaWorldEnv",
    kwargs={
        "env_name": "drawer-close-v2",
        "max_steps": 500,
    },
)


envs.register(
    id="drawer-open-v2",
    entry_point="diffgro.environments.metaworld.envs.env:MetaWorldEnv",
    kwargs={
        "env_name": "drawer-open-v2",
        "max_steps": 500,
    },
)


envs.register(
    id="peg-insert-side-v2",
    entry_point="diffgro.environments.metaworld.envs.env:MetaWorldEnv",
    kwargs={
        "env_name": "peg-insert-side-v2",
        "max_steps": 500,
    },
)


envs.register(
    id="faucet-open-v2",
    entry_point="diffgro.environments.metaworld.envs.env:MetaWorldEnv",
    kwargs={
        "env_name": "faucet-open-v2",
        "max_steps": 500,
    },
)

envs.register(
    id="push-v2",
    entry_point="diffgro.environments.metaworld.envs.env:MetaWorldEnv",
    kwargs={
        "env_name": "push-v2",
        "max_steps": 500,
    },
)

envs.register(
    id="reach-v2",
    entry_point="diffgro.environments.metaworld.envs.env:MetaWorldEnv",
    kwargs={
        "env_name": "reach-v2",
        "max_steps": 500,
    },
)

envs.register(
    id="window-open-v2",
    entry_point="diffgro.environments.metaworld.envs.env:MetaWorldEnv",
    kwargs={
        "env_name": "window-open-v2",
        "max_steps": 500,
    },
)

envs.register(
    id="window-close-v2",
    entry_point="diffgro.environments.metaworld.envs.env:MetaWorldEnv",
    kwargs={
        "env_name": "window-close-v2",
        "max_steps": 500,
    },
)

########## Variant Environments ##########

envs.register(
    id="button-press-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:ButtonPressVariantEnv",
    kwargs=button_press_variant_config,
)

envs.register(
    id="door-open-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:DoorOpenVariantEnv",
    kwargs=door_open_variant_config,
)


envs.register(
    id="faucet-open-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:FaucetOpenVariantEnv",
    kwargs=faucet_open_variant_config,
)

envs.register(
    id="drawer-open-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:DrawerVariantEnv",
    kwargs=drawer_open_variant_config,
)


envs.register(
    id="drawer-close-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:DrawerVariantEnv",
    kwargs=drawer_close_variant_config,
)

envs.register(
    id="window-open-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:WindowVariantEnv",
    kwargs=window_open_variant_config,
)


envs.register(
    id="window-close-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:WindowVariantEnv",
    kwargs=window_close_variant_config,
)

envs.register(
    id="reach-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:ReachVariantEnv",
    kwargs=reach_variant_config,
)

envs.register(
    id="push-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:PushVariantEnv",
    kwargs=push_variant_config,
)

envs.register(
    id="peg-insert-side-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:PegInsertSideVariantEnv",
    kwargs=peg_insert_side_variant_config,
)


####################################################################################


envs.register(
    id="push-back-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:PushBackVariantEnv",
    kwargs=push_back_variant_config,
)

envs.register(
    id="plate-slide-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:PlateSlideVariantEnv",
    kwargs=plate_slide_variant_config,
)

envs.register(
    id="plate-slide-side-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:PlateSlideSideVariantEnv",
    kwargs=plate_slide_side_config,
)

envs.register(
    id="plate-slide-back-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:PlateSlideBackVariantEnv",
    kwargs=plate_slide_back_variant_config,
)

envs.register(
    id="plate-slide-back-side-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:PlateSlideBackSideVariantEnv",
    kwargs=plate_slide_back_side_variant_config,
)

envs.register(
    id="peg-unplug-side-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:PegUnplugSideVariantEnv",
    kwargs=peg_unplug_side_variant_config,
)

envs.register(
    id="pick-place-variant-v2",
    entry_point="diffgro.environments.metaworld.envs:PickPlaceVariantEnv",
    kwargs=pick_place_variant_config,
)
