from environments.dnc_environments.picker import PickerEnv

# from environments.dnc_environments.lob import LobberEnv
# from environments.dnc_environments.ant import AntEnv
# from environments.dnc_environments.catch import CatchEnv

from environments.dnc_environments.base import create_env_partitions

import numpy as np

_envs = {"pick": PickerEnv}  # , "lob": LobberEnv, "catch": CatchEnv, "ant": AntEnv}

_stochastic_params = {
    "pick": dict(goal_args=("noisy", (0.6, 0.2), 0.1)),
    "lob": dict(box_center=(0, 0), box_noise=0.4),
    "catch": dict(start_pos=(0.1, 1.7), start_noise=0.2),
    "ant": dict(angle_range=(0, 2 * np.pi)),
}

_deterministic_params = {
    "pick": dict(goal_args=("noisy", (0.6, 0.2), 0)),
    "lob": dict(box_center=(0, 0), box_noise=0),
    "catch": dict(start_pos=(0.1, 1.7), start_noise=0),
    "ant": dict(angle_range=(-1e-4, 1e-4)),
}


def create_stochastic(name):
    assert name in _stochastic_params
    return _envs[name](**_stochastic_params[name])


def create_deterministic(name):
    assert name in _deterministic_params
    return _envs[name](**_deterministic_params[name])
