from gym.envs.registration import register
from d4rl_ext.locomotion import ant
from d4rl_ext.locomotion import maze_env

register(
    id='antmaze-ultra-diverse-v0',
    entry_point='d4rl_ext.locomotion.ant:make_ant_maze_env',
    max_episode_steps=2000,
    kwargs={
        'deprecated': True,
        'maze_map': maze_env.ULTRA_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':False,
        'dataset_url':'https://github.com/ZhengyaoJiang/d4rl/releases/download/public/Ant_maze_ultra_noisy_multistart_True_multigoal_True_sparse.hdf5',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)

register(
    id='antmaze-ultra-play-v0',
    entry_point='d4rl_ext.locomotion.ant:make_ant_maze_env',
    max_episode_steps=2000,
    kwargs={
        'deprecated': True,
        'maze_map': maze_env.ULTRA_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':False,
        'dataset_url':'https://github.com/ZhengyaoJiang/d4rl/releases/download/public/Ant_maze_ultra_noisy_multistart_True_multigoal_False_sparse.hdf5',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)

register(
    id='antmaze-extreme-play-v0',
    entry_point='d4rl_ext.locomotion.ant:make_ant_maze_env',
    max_episode_steps=3000,
    kwargs={
        'deprecated': True,
        'maze_map': maze_env.EXTREME_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':False,
        'dataset_url':'https://NO_WEBSITE_FOR_NOW/Ant_maze_extreme-maze_noisy_multistart_True_multigoal_False.hdf5/raw?ref=main&private_token=97NvzKm_cG9Xz7ADZWPe&play',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)

register(
    id='antmaze-extreme-diverse-v0',
    entry_point='d4rl_ext.locomotion.ant:make_ant_maze_env',
    max_episode_steps=3000,
    kwargs={
        'deprecated': True,
        'maze_map': maze_env.EXTREME_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':False,
        'dataset_url':'https://NO_WEBSITE_FOR_NOW/Ant_maze_extreme-maze_noisy_multistart_True_multigoal_True.hdf5/raw?ref=main&private_token=97NvzKm_cG9Xz7ADZWPe&diverse',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)


####################
# random init envs #
####################
register(
    id='random-antmaze-medium-play-v2',
    entry_point='d4rl_ext.locomotion.ant:make_random_ant_maze_env',
    max_episode_steps=1000,
    kwargs={
        'maze_map': maze_env.BIG_MAZE_TEST,
        'reward_type':'sparse',
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_big-maze_noisy_multistart_True_multigoal_False_sparse_fixed.hdf5',
        'non_zero_reset':True, 
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)

register(
    id='random-antmaze-medium-diverse-v2',
    entry_point='d4rl_ext.locomotion.ant:make_random_ant_maze_env',
    max_episode_steps=1000,
    kwargs={
        'maze_map': maze_env.BIG_MAZE_TEST,
        'reward_type':'sparse',
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_big-maze_noisy_multistart_True_multigoal_True_sparse_fixed.hdf5',
        'non_zero_reset':True, 
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)

register(
    id='random-antmaze-large-diverse-v2',
    entry_point='d4rl_ext.locomotion.ant:make_random_ant_maze_env',
    max_episode_steps=1000,
    kwargs={
        'maze_map': maze_env.HARDEST_MAZE_TEST,
        'reward_type':'sparse',
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_hardest-maze_noisy_multistart_True_multigoal_True_sparse_fixed.hdf5',
        'non_zero_reset':True, 
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)

register(
    id='random-antmaze-large-play-v2',
    entry_point='d4rl_ext.locomotion.ant:make_random_ant_maze_env',
    max_episode_steps=1000,
    kwargs={
        'maze_map': maze_env.HARDEST_MAZE_TEST,
        'reward_type':'sparse',
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_hardest-maze_noisy_multistart_True_multigoal_False_sparse_fixed.hdf5',
        'non_zero_reset':True, 
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)

register(
    id='random-antmaze-ultra-diverse-v0',
    entry_point='d4rl_ext.locomotion.ant:make_random_ant_maze_env',
    max_episode_steps=2000,
    kwargs={
        'deprecated': True,
        'maze_map': maze_env.ULTRA_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':True,
        'dataset_url':'https://github.com/ZhengyaoJiang/d4rl/releases/download/public/Ant_maze_ultra_noisy_multistart_True_multigoal_True_sparse.hdf5',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)

register(
    id='random-antmaze-ultra-play-v0',
    entry_point='d4rl_ext.locomotion.ant:make_random_ant_maze_env',
    max_episode_steps=2000,
    kwargs={
        'deprecated': True,
        'maze_map': maze_env.ULTRA_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':True,
        'dataset_url':'https://github.com/ZhengyaoJiang/d4rl/releases/download/public/Ant_maze_ultra_noisy_multistart_True_multigoal_True_sparse.hdf5',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)

register(
    id='random-antmaze-extreme-play-v0',
    entry_point='d4rl_ext.locomotion.ant:make_random_ant_maze_env',
    max_episode_steps=3000,
    kwargs={
        'deprecated': True,
        'maze_map': maze_env.EXTREME_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':True,
        'dataset_url':'https://NO_WEBSITE_FOR_NOW/Ant_maze_extreme-maze_noisy_multistart_True_multigoal_False.hdf5/raw?ref=main&private_token=97NvzKm_cG9Xz7ADZWPe&play',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)

register(
    id='random-antmaze-extreme-diverse-v0',
    entry_point='d4rl_ext.locomotion.ant:make_random_ant_maze_env',
    max_episode_steps=3000,
    kwargs={
        'deprecated': True,
        'maze_map': maze_env.EXTREME_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':True,
        'dataset_url':'https://NO_WEBSITE_FOR_NOW/Ant_maze_extreme-maze_noisy_multistart_True_multigoal_True.hdf5/raw?ref=main&private_token=97NvzKm_cG9Xz7ADZWPe&diverse',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
    }
)

#################
# teleport envs #
#################
register(
    id='teleport-antmaze-medium-play-v2',
    entry_point='d4rl_ext.locomotion.ant:make_teleport_ant_maze_env',
    max_episode_steps=1000,
    kwargs={
        'maze_map': maze_env.BIG_MAZE_TEST,
        'reward_type':'sparse',
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_big-maze_noisy_multistart_True_multigoal_False_sparse_fixed.hdf5',
        'non_zero_reset':False, 
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
        'p_teleport': 10/1000,
    }
)

register(
    id='teleport-antmaze-medium-diverse-v2',
    entry_point='d4rl_ext.locomotion.ant:make_teleport_ant_maze_env',
    max_episode_steps=1000,
    kwargs={
        'maze_map': maze_env.BIG_MAZE_TEST,
        'reward_type':'sparse',
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_big-maze_noisy_multistart_True_multigoal_True_sparse_fixed.hdf5',
        'non_zero_reset':False, 
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
        'p_teleport': 10/1000,
    }
)

register(
    id='teleport-antmaze-large-diverse-v2',
    entry_point='d4rl_ext.locomotion.ant:make_teleport_ant_maze_env',
    max_episode_steps=1000,
    kwargs={
        'maze_map': maze_env.HARDEST_MAZE_TEST,
        'reward_type':'sparse',
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_hardest-maze_noisy_multistart_True_multigoal_True_sparse_fixed.hdf5',
        'non_zero_reset':False, 
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
        'p_teleport': 10/1000,
    }
)

register(
    id='teleport-antmaze-large-play-v2',
    entry_point='d4rl_ext.locomotion.ant:make_teleport_ant_maze_env',
    max_episode_steps=1000,
    kwargs={
        'maze_map': maze_env.HARDEST_MAZE_TEST,
        'reward_type':'sparse',
        'dataset_url':'http://rail.eecs.berkeley.edu/datasets/offline_rl/ant_maze_v2/Ant_maze_hardest-maze_noisy_multistart_True_multigoal_False_sparse_fixed.hdf5',
        'non_zero_reset':False, 
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
        'p_teleport': 10/1000,
    }
)

register(
    id='teleport-antmaze-ultra-diverse-v0',
    entry_point='d4rl_ext.locomotion.ant:make_teleport_ant_maze_env',
    max_episode_steps=2000,
    kwargs={
        'deprecated': True,
        'maze_map': maze_env.ULTRA_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':False,
        'dataset_url':'https://github.com/ZhengyaoJiang/d4rl/releases/download/public/Ant_maze_ultra_noisy_multistart_True_multigoal_True_sparse.hdf5',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
        'p_teleport': 10/2000,
    }
)

register(
    id='teleport-antmaze-ultra-play-v0',
    entry_point='d4rl_ext.locomotion.ant:make_teleport_ant_maze_env',
    max_episode_steps=2000,
    kwargs={
        'deprecated': True,
        'maze_map': maze_env.ULTRA_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':False,
        'dataset_url':'https://github.com/ZhengyaoJiang/d4rl/releases/download/public/Ant_maze_ultra_noisy_multistart_True_multigoal_True_sparse.hdf5',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
        'p_teleport': 10/2000,
    }
)

register(
    id='teleport-antmaze-extreme-play-v0',
    entry_point='d4rl_ext.locomotion.ant:make_teleport_ant_maze_env',
    max_episode_steps=3000,
    kwargs={
        'deprecated': True,
        'maze_map': maze_env.EXTREME_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':False,
        'dataset_url':'https://NO_WEBSITE_FOR_NOW/Ant_maze_extreme-maze_noisy_multistart_True_multigoal_False.hdf5/raw?ref=main&private_token=97NvzKm_cG9Xz7ADZWPe&play',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
        'p_teleport': 10/3000,
    }
)

register(
    id='teleport-antmaze-extreme-diverse-v0',
    entry_point='d4rl_ext.locomotion.ant:make_teleport_ant_maze_env',
    max_episode_steps=3000,
    kwargs={
        'deprecated': False,
        'maze_map': maze_env.EXTREME_MAZE_TEST,
        'reward_type':'sparse',
        'non_zero_reset':False,
        'dataset_url':'https://NO_WEBSITE_FOR_NOW/Ant_maze_extreme-maze_noisy_multistart_True_multigoal_True.hdf5/raw?ref=main&private_token=97NvzKm_cG9Xz7ADZWPe&diverse',
        'eval':True,
        'maze_size_scaling': 4.0,
        'ref_min_score': 0.0,
        'ref_max_score': 1.0,
        'p_teleport': 10/3000,
    }
)