from diffgro.environments.metaworld.policies.policy import *


metaworld_exp_dict = {
    "button-press-v2": SkillButtonPressV2Policy,
    "door-open-v2": SkillDoorOpenV2Policy,
    "drawer-close-v2": SkillDrawerCloseV2Policy,
    "drawer-open-v2": SkillDrawerOpenV2Policy,
    "peg-insert-side-v2": SkillPegInsertSideV2Policy,
    "push-v2": SkillPushV2Policy,
    "reach-v2": SkillReachV2Policy,
    "window-open-v2": SkillWindowOpenV2Policy,
    "window-close-v2": SkillWindowCloseV2Policy,
    "faucet-open-v2": SkillFaucetOpenV2Policy,
    #############################################################
    "push-back-v2": SkillPushBackV2Policy,
    "plate-slide-v2": SkillPlateSlideV2Policy,
    "plate-slide-side-v2": SkillPlateSlideSideV2Policy,
    "plate-slide-back-v2": SkillPlateSlideSideV2Policy,
    "plate-slide-back-side-v2": SkillPlateSlideSideV2Policy,
    "peg-unplug-side-v2": SkillPegUnplugSideV2Policy,
    # "pick-place-v2": SkillPickPlaceV2Policy,
    "pick-place-wall-v2": SkillPickPlaceV2Policy,  # alias
}

variant_metaworld_exp_dict = {
    "button-press-v2": SkillButtonPressVariantV2Policy,
    "door-open-v2": SkillDoorOpenVariantV2Policy,
    "drawer-close-v2": SkillDrawerCloseVariantV2Policy,
    "drawer-open-v2": SkillDrawerOpenVariantV2Policy,
    "peg-insert-side-v2": SkillPegInsertSideVariantV2Policy,
    # "push-v2": SkillPushVariantV2Policy,
    "push-wall-v2": SkillPushVariantV2Policy,  # alias
    # "reach-v2": SkillReachVariantV2Policy,
    "reach-wall-v2": SkillReachVariantV2Policy,  # alias
    "window-open-v2": SkillWindowOpenVariantV2Policy,
    "window-close-v2": SkillWindowCloseVariantV2Policy,
    "faucet-open-v2": SkillFaucetOpenVariantV2Policy,
    #############################################################
    "push-back-v2": SkillPushBackVariantV2Policy,
    "plate-slide-v2": SkillPlateSlideVariantV2Policy,
    "plate-slide-side-v2": SkillPlateSlideSideVariantV2Policy,
    "plate-slide-back-v2": SkillPlateSlideBackVariantV2Policy,
    "plate-slide-back-side-v2": SkillPlateSlideBackSideVariantV2Policy,
    "peg-unplug-side-v2": SkillPegUnplugSideVariantV2Policy,
    # "pick-place-v2": SkillPickPlaceVariantV2Policy,
    "pick-place-wall-v2": SkillPickPlaceVariantV2Policy,  # alias
}


def make_metaworld_policy(env_name, variant=None):
    if variant is None:
        return metaworld_exp_dict[env_name]()
    else:
        return variant_metaworld_exp_dict[env_name](variant)
