from diffgro.environments.variant import *


door_open_variant_config = {
    "env_name": "door-open-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            # "handle_position": Uniform(low=[-0.03, 0, -0.02], high=[0, 0, 0.02]),
            "handle_position": Uniform(low=[0, 0, 0], high=[0, 0, 0]),
            "goal_resistance": Categorical(a=[0]),
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}


button_press_variant_config = {
    "env_name": "button-press-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            # "button_depth": Uniform(low=-0.015, high=0.02),
            "button_depth": Uniform(low=0, high=0),
            "goal_resistance": Categorical(a=[0]),
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

drawer_open_variant_config = {
    "env_name": "drawer-open-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "handle_position": Uniform(low=[0, 0, 0], high=[0, 0, 0]),
            "goal_resistance": Categorical(a=[0]),
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

drawer_close_variant_config = {
    "env_name": "drawer-close-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            # "handle_position": Uniform(low=[0, 0, -0.015], high=[0, 0, 0.015]),
            "handle_position": Uniform(low=[0, 0, 0], high=[0, 0, 0]),
            "goal_resistance": Categorical(a=[0]),
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

window_open_variant_config = {
    "env_name": "window-open-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            # "handle_position": Uniform(low=[0, 0, -0.03], high=[0, 0, 0.03]),
            "handle_position": Uniform(low=[0, 0, 0], high=[0, 0, 0]),
            "goal_resistance": Categorical(a=[0]),
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

window_close_variant_config = {
    "env_name": "window-close-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "handle_position": Uniform(low=[0, 0, 0], high=[0, 0, 0]),
            "goal_resistance": Categorical(a=[0]),
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

faucet_open_variant_config = {
    "env_name": "faucet-open-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "initial_angle": Uniform(low=0, high=0),
            "goal_resistance": Categorical(a=[0]),
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

reach_variant_config = {
    "env_name": "reach-wall-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "wall_existence": Categorical(a=[False]),
            "goal_resistance": Categorical(a=[0]),  # dummy
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

push_variant_config = {
    "env_name": "push-wall-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "wall_existence": Categorical(a=[False]),
            "goal_resistance": Categorical(a=[0]),  # dummy
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

peg_insert_side_variant_config = {
    "env_name": "peg-insert-side-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "peg_shape": Normal(loc=[0.015, 0.015, 0.12], scale=[0, 0, 0]),
            "goal_resistance": Categorical(a=[0]),  # dummy
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

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

push_back_variant_config = {
    "env_name": "push-back-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "goal_resistance": Categorical(a=[0]),  # dummy
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

plate_slide_variant_config = {
    "env_name": "plate-slide-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "goal_resistance": Categorical(a=[0]),  # dummy
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

plate_slide_side_config = {
    "env_name": "plate-slide-side-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "goal_resistance": Categorical(a=[0]),  # dummy
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}


plate_slide_back_variant_config = {
    "env_name": "plate-slide-back-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "goal_resistance": Categorical(a=[0]),  # dummy
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}


plate_slide_back_side_variant_config = {
    "env_name": "plate-slide-back-side-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "goal_resistance": Categorical(a=[0]),  # dummy
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}

peg_unplug_side_variant_config = {
    "env_name": "peg-unplug-side-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "goal_resistance": Categorical(a=[0]),  # dummy
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}


pick_place_variant_config = {
    "env_name": "pick-place-wall-v2",
    "max_steps": 500,
    "variant_space": VariantSpace(
        {
            "arm_speed": Categorical(a=[1.0]),
            "wall_existence": Categorical(a=[False]),
            "goal_resistance": Categorical(a=[0]),  # dummy
            "wind_xspeed": Categorical(a=[1.0]),
            "wind_yspeed": Categorical(a=[1.0]),
        }
    ),
}
