from gym.envs.registration import register

from alr_envs.stochastic_search.functions.f_rosenbrock import Rosenbrock

register(
    id='ALRReacher-v0',
    entry_point='alr_envs.mujoco:ALRReacherEnv',
    max_episode_steps=200,
    kwargs={
        "steps_before_reward": 0,
        "n_links": 5,
    }
)

register(
    id='ALRReacherSparse-v0',
    entry_point='alr_envs.mujoco:ALRReacherEnv',
    max_episode_steps=200,
    kwargs={
        "steps_before_reward": 200,
        "n_links": 5,
    }
)

register(
    id='ALRReacherSparseBalanced-v0',
    entry_point='alr_envs.mujoco:ALRReacherEnv',
    max_episode_steps=200,
    kwargs={
        "steps_before_reward": 200,
        "n_links": 5,
        "balance": True,
    }
)

register(
    id='ALRReacherShort-v0',
    entry_point='alr_envs.mujoco:ALRReacherEnv',
    max_episode_steps=50,
    kwargs={
        "steps_before_reward": 0,
        "n_links": 5,
    }
)

register(
    id='ALRReacherShortSparse-v0',
    entry_point='alr_envs.mujoco:ALRReacherEnv',
    max_episode_steps=50,
    kwargs={
        "steps_before_reward": 50,
        "n_links": 5,
    }
)

register(
    id='ALRReacher7-v0',
    entry_point='alr_envs.mujoco:ALRReacherEnv',
    max_episode_steps=200,
    kwargs={
        "steps_before_reward": 0,
        "n_links": 7,
    }
)

# register(
#     id='ALRReacherSparse-v0',
#     entry_point='alr_envs.mujoco:ALRReacherEnv',
#     max_episode_steps=200,
#     kwargs={
#         "steps_before_reward": 200,
#         "n_links": 7,
#     }
# )

register(
    id='ALRReacher7Short-v0',
    entry_point='alr_envs.mujoco:ALRReacherEnv',
    max_episode_steps=50,
    kwargs={
        "steps_before_reward": 0,
        "n_links": 7,
    }
)

register(
    id='ALRReacher7ShortSparse-v0',
    entry_point='alr_envs.mujoco:ALRReacherEnv',
    max_episode_steps=50,
    kwargs={
        "steps_before_reward": 50,
        "n_links": 7,
    }
)

register(
    id='SimpleReacher-v0',
    entry_point='alr_envs.classic_control:SimpleReacherEnv',
    max_episode_steps=200,
    kwargs={
        "n_links": 2,
    }
)

register(
    id='SimpleReacher5-v0',
    entry_point='alr_envs.classic_control:SimpleReacherEnv',
    max_episode_steps=200,
    kwargs={
        "n_links": 5,
    }
)

for dim in [5, 10, 25, 50, 100]:
    register(
        id=f'Rosenbrock{dim}-v0',
        entry_point='alr_envs.stochastic_search:StochasticSearchEnv',
        max_episode_steps=1,
        kwargs={
            "cost_f": Rosenbrock,
        }
    )