# Copyright 2017 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.
# ============================================================================

"""Container classes used in control domains."""

import collections

_NAME_ALREADY_EXISTS = (
    "A function named {name!r} already exists in the container and "
    "`allow_overriding_keys` is False.")


class TaggedTasks(collections.abc.Mapping):
  """Maps task names to their corresponding factory functions with tags.

  To store a function in a `TaggedTasks` container, we can use its `.add`
  decorator:

  ```python
  tasks = TaggedTasks()

  @tasks.add('easy', 'stable')
  def example_task():
    ...
    return environment

  environment_factory = tasks['example_task']

  # Or to restrict to a given tag:
  environment_factory = tasks.tagged('easy')['example_task']
  ```
  """

  def __init__(self, allow_overriding_keys=False):
    """Initializes a new `TaggedTasks` container.

    Args:
      allow_overriding_keys: Boolean, whether `add` can override existing keys
        within the container. If False (default), calling `add` multiple times
        with the same function name will result in a `ValueError`.
    """
    self._tasks = collections.OrderedDict()
    self._tags = collections.defaultdict(dict)
    self.allow_overriding_keys = allow_overriding_keys

  def add(self, *tags):
    """Decorator that adds a factory function to the container with tags.

    Args:
      *tags: Strings specifying the tags for this function.

    Returns:
      The same function.

    Raises:
      ValueError: if a function with the same name already exists within the
        container and `allow_overriding_keys` is False.
    """
    def wrap(factory_func):
      name = factory_func.__name__
      if name in self and not self.allow_overriding_keys:
        raise ValueError(_NAME_ALREADY_EXISTS.format(name=name))
      self._tasks[name] = factory_func
      for tag in tags:
        self._tags[tag][name] = factory_func
      return factory_func
    return wrap

  def tagged(self, *tags):
    """Returns a (possibly empty) dict of functions matching all the given tags.

    Args:
      *tags: Strings specifying tags to query by.

    Returns:
      A dict of `{name: function}` containing all the functions that are tagged
      by all of the strings in `tags`.
    """
    if not tags:
      return {}
    tags = set(tags)
    if not tags.issubset(self._tags.keys()):
      return {}
    names = self._tags[tags.pop()].keys()
    while tags:
      names &= self._tags[tags.pop()].keys()
    return {name: self._tasks[name] for name in names}

  def tags(self):
    """Returns a list of all the tags in this container."""
    return list(self._tags.keys())

  def __getitem__(self, k):
    return self._tasks[k]

  def __iter__(self):
    return iter(self._tasks)

  def __len__(self):
    return len(self._tasks)

  def __repr__(self):
    return "{}({})".format(self.__class__.__name__, str(self._tasks))