# Copyright 2018 The dm_control Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

"""A structured set of manipulation tasks with a single entry point."""


from absl import flags
from dm_control import composer as _composer
from dm_control.manipulation import bricks as _bricks
from dm_control.manipulation import lift as _lift
from dm_control.manipulation import place as _place
from dm_control.manipulation import reach as _reach
from dm_control.manipulation.shared import registry as _registry

_registry.done_importing_tasks()

_TIME_LIMIT = 10.
_TIMEOUT = None

ALL = tuple(_registry.get_all_names())
TAGS = tuple(_registry.get_tags())

flags.DEFINE_bool('timeout', True, 'Whether episodes should have a time limit.')
FLAGS = flags.FLAGS


def _get_timeout():
  global _TIMEOUT
  if _TIMEOUT is None:
    if FLAGS.is_parsed():
      _TIMEOUT = FLAGS.timeout
    else:
      _TIMEOUT = FLAGS['timeout'].default
  return _TIMEOUT


def get_environments_by_tag(tag):
  """Returns the names of all environments matching a given tag.

  Args:
    tag: A string from `TAGS`.

  Returns:
    A tuple of environment names.
  """
  return tuple(_registry.get_names_by_tag(tag))


def load(environment_name, seed=None):
  """Loads a manipulation environment.

  Args:
    environment_name: String, the name of the environment to load. Must be in
      `ALL`.
    seed: An optional integer used to seed the task's random number generator.
      If None (default), the random number generator will self-seed from a
      platform-dependent source of entropy.

  Returns:
    An instance of `composer.Environment`.
  """
  task = _registry.get_constructor(environment_name)()
  time_limit = _TIME_LIMIT if _get_timeout() else float('inf')
  return _composer.Environment(task, time_limit=time_limit, random_state=seed)
