# pylint: disable=g-bad-file-header
# Copyright 2019 The dm_env Authors. All Rights Reserved.
#
# 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.
# ============================================================================

"""Classes that describe numpy arrays."""

import inspect
from typing import Optional

import numpy as np

_INVALID_SHAPE = 'Expected shape %r but found %r'
_INVALID_DTYPE = 'Expected dtype %r but found %r'
_OUT_OF_BOUNDS = 'Values were not all within bounds %s <= %s <= %s'
_VAR_ARGS_NOT_ALLOWED = 'Spec subclasses must not accept *args.'
_VAR_KWARGS_NOT_ALLOWED = 'Spec subclasses must not accept **kwargs.'
_MINIMUM_MUST_BE_LESS_THAN_OR_EQUAL_TO_MAXIMUM = (
    'All values in `minimum` must be less than or equal to their corresponding '
    'value in `maximum`, got:\nminimum={minimum!r}\nmaximum={maximum!r}.')
_MINIMUM_INCOMPATIBLE_WITH_SHAPE = '`minimum` is incompatible with `shape`'
_MAXIMUM_INCOMPATIBLE_WITH_SHAPE = '`maximum` is incompatible with `shape`'


class Array:
  """Describes a numpy array or scalar shape and dtype.

  An `Array` spec allows an API to describe the arrays that it accepts or
  returns, before that array exists.
  The equivalent version describing a `tf.Tensor` is `TensorSpec`.
  """
  __slots__ = ('_shape', '_dtype', '_name')
  __hash__ = None

  def __init__(self, shape, dtype, name: Optional[str] = None):
    """Initializes a new `Array` spec.

    Args:
      shape: An iterable specifying the array shape.
      dtype: numpy dtype or string specifying the array dtype.
      name: Optional string containing a semantic name for the corresponding
        array. Defaults to `None`.

    Raises:
      TypeError: If `shape` is not an iterable of elements convertible to int,
      or if `dtype` is not convertible to a numpy dtype.
    """
    self._shape = tuple(int(dim) for dim in shape)
    self._dtype = np.dtype(dtype)
    self._name = name

  @property
  def shape(self):
    """Returns a `tuple` specifying the array shape."""
    return self._shape

  @property
  def dtype(self):
    """Returns a numpy dtype specifying the array dtype."""
    return self._dtype

  @property
  def name(self):
    """Returns the name of the Array."""
    return self._name

  def __repr__(self):
    return 'Array(shape={}, dtype={}, name={})'.format(self.shape,
                                                       repr(self.dtype),
                                                       repr(self.name))

  def __eq__(self, other):
    """Checks if the shape and dtype of two specs are equal."""
    if not isinstance(other, Array):
      return False
    return self.shape == other.shape and self.dtype == other.dtype

  def __ne__(self, other):
    return not self == other

  def _fail_validation(self, message, *args):
    message %= args
    if self.name:
      message += ' for spec %s' % self.name
    raise ValueError(message)

  def validate(self, value):
    """Checks if value conforms to this spec.

    Args:
      value: a numpy array or value convertible to one via `np.asarray`.

    Returns:
      value, converted if necessary to a numpy array.

    Raises:
      ValueError: if value doesn't conform to this spec.
    """
    value = np.asarray(value)
    if value.shape != self.shape:
      self._fail_validation(_INVALID_SHAPE, self.shape, value.shape)
    if value.dtype != self.dtype:
      self._fail_validation(_INVALID_DTYPE, self.dtype, value.dtype)
    return value

  def generate_value(self):
    """Generate a test value which conforms to this spec."""
    return np.zeros(shape=self.shape, dtype=self.dtype)

  def _get_constructor_kwargs(self):
    """Returns constructor kwargs for instantiating a new copy of this spec."""
    # Get the names and kinds of the constructor parameters.
    params = inspect.signature(type(self)).parameters
    # __init__ must not accept *args or **kwargs, since otherwise we won't be
    # able to infer what the corresponding attribute names are.
    kinds = {value.kind for value in params.values()}
    if inspect.Parameter.VAR_POSITIONAL in kinds:
      raise TypeError(_VAR_ARGS_NOT_ALLOWED)
    elif inspect.Parameter.VAR_KEYWORD in kinds:
      raise TypeError(_VAR_KWARGS_NOT_ALLOWED)
    # Note that we assume direct correspondence between the names of constructor
    # arguments and attributes.
    return {name: getattr(self, name) for name in params.keys()}

  def replace(self, **kwargs):
    """Returns a new copy of `self` with specified attributes replaced.

    Args:
      **kwargs: Optional attributes to replace.

    Returns:
      A new copy of `self`.
    """
    all_kwargs = self._get_constructor_kwargs()
    all_kwargs.update(kwargs)
    return type(self)(**all_kwargs)

  def __reduce__(self):
    return Array, (self._shape, self._dtype, self._name)


class BoundedArray(Array):
  """An `Array` spec that specifies minimum and maximum values.

  Example usage:
  ```python
  # Specifying the same minimum and maximum for every element.
  spec = BoundedArray((3, 4), np.float64, minimum=0.0, maximum=1.0)

  # Specifying a different minimum and maximum for each element.
  spec = BoundedArray(
      (2,), np.float64, minimum=[0.1, 0.2], maximum=[0.9, 0.9])

  # Specifying the same minimum and a different maximum for each element.
  spec = BoundedArray(
      (3,), np.float64, minimum=-10.0, maximum=[4.0, 5.0, 3.0])
  ```

  Bounds are meant to be inclusive. This is especially important for
  integer types. The following spec will be satisfied by arrays
  with values in the set {0, 1, 2}:
  ```python
  spec = BoundedArray((3, 4), int, minimum=0, maximum=2)
  ```

  Note that one or both bounds may be infinite. For example, the set of
  non-negative floats can be expressed as:
  ```python
  spec = BoundedArray((), np.float64, minimum=0.0, maximum=np.inf)
  ```
  In this case `np.inf` would be considered valid, since the upper bound is
  inclusive.
  """
  __slots__ = ('_minimum', '_maximum')
  __hash__ = None

  def __init__(self, shape, dtype, minimum, maximum, name=None):
    """Initializes a new `BoundedArray` spec.

    Args:
      shape: An iterable specifying the array shape.
      dtype: numpy dtype or string specifying the array dtype.
      minimum: Number or sequence specifying the minimum element bounds
        (inclusive). Must be broadcastable to `shape`.
      maximum: Number or sequence specifying the maximum element bounds
        (inclusive). Must be broadcastable to `shape`.
      name: Optional string containing a semantic name for the corresponding
        array. Defaults to `None`.

    Raises:
      ValueError: If `minimum` or `maximum` are not broadcastable to `shape`.
      ValueError: If any values in `minimum` are greater than their
        corresponding value in `maximum`.
      TypeError: If the shape is not an iterable or if the `dtype` is an invalid
        numpy dtype.
    """
    super(BoundedArray, self).__init__(shape, dtype, name)

    try:
      bcast_minimum = np.broadcast_to(minimum, shape=shape)
    except ValueError as numpy_exception:
      raise ValueError(_MINIMUM_INCOMPATIBLE_WITH_SHAPE) from numpy_exception
    try:
      bcast_maximum = np.broadcast_to(maximum, shape=shape)
    except ValueError as numpy_exception:
      raise ValueError(_MAXIMUM_INCOMPATIBLE_WITH_SHAPE) from numpy_exception

    if np.any(bcast_minimum > bcast_maximum):
      raise ValueError(_MINIMUM_MUST_BE_LESS_THAN_OR_EQUAL_TO_MAXIMUM.format(
          minimum=minimum, maximum=maximum))

    self._minimum = np.array(minimum, dtype=self.dtype)
    self._minimum.setflags(write=False)

    self._maximum = np.array(maximum, dtype=self.dtype)
    self._maximum.setflags(write=False)

  @property
  def minimum(self):
    """Returns a NumPy array specifying the minimum bounds (inclusive)."""
    return self._minimum

  @property
  def maximum(self):
    """Returns a NumPy array specifying the maximum bounds (inclusive)."""
    return self._maximum

  def __repr__(self):
    template = ('BoundedArray(shape={}, dtype={}, name={}, '
                'minimum={}, maximum={})')
    return template.format(self.shape, repr(self.dtype), repr(self.name),
                           self._minimum, self._maximum)

  def __eq__(self, other):
    if not isinstance(other, BoundedArray):
      return False
    return (super(BoundedArray, self).__eq__(other) and
            (self.minimum == other.minimum).all() and
            (self.maximum == other.maximum).all())

  def validate(self, value):
    value = np.asarray(value)
    super(BoundedArray, self).validate(value)
    if (value < self.minimum).any() or (value > self.maximum).any():
      self._fail_validation(_OUT_OF_BOUNDS, self.minimum, value, self.maximum)
    return value

  def generate_value(self):
    return (np.ones(shape=self.shape, dtype=self.dtype) *
            self.dtype.type(self.minimum))

  def __reduce__(self):
    return BoundedArray, (self._shape, self._dtype, self._minimum,
                          self._maximum, self._name)


_NUM_VALUES_NOT_POSITIVE = '`num_values` must be a positive integer, got {}.'
_DTYPE_NOT_INTEGRAL = '`dtype` must be integral, got {}.'
_DTYPE_OVERFLOW = (
    '`dtype` {} is not big enough to hold `num_values` ({}) without overflow.')


class DiscreteArray(BoundedArray):
  """Represents a discrete, scalar, zero-based space.

  This is a special case of the parent `BoundedArray` class. It represents a
  0-dimensional numpy array  containing a single integer value between
  0 and num_values - 1 (inclusive), and exposes a scalar `num_values` property
  in addition to the standard `BoundedArray` interface.

  For an example use-case, this can be used to define the action space of a
  simple RL environment that accepts discrete actions.
  """

  _REPR_TEMPLATE = (
      'DiscreteArray(shape={self.shape}, dtype={self.dtype}, name={self.name}, '
      'minimum={self.minimum}, maximum={self.maximum}, '
      'num_values={self.num_values})')

  __slots__ = ('_num_values',)

  def __init__(self, num_values, dtype=np.int32, name=None):
    """Initializes a new `DiscreteArray` spec.

    Args:
      num_values: Integer specifying the number of possible values to represent.
      dtype: The dtype of the array. Must be an integral type large enough to
        hold `num_values` without overflow.
      name: Optional string specifying the name of the array.

    Raises:
      ValueError: If `num_values` is not positive, if `dtype` is not integral,
        or if `dtype` is not large enough to hold `num_values` without overflow.
    """
    if num_values <= 0 or not np.issubdtype(type(num_values), np.integer):
      raise ValueError(_NUM_VALUES_NOT_POSITIVE.format(num_values))

    if not np.issubdtype(dtype, np.integer):
      raise ValueError(_DTYPE_NOT_INTEGRAL.format(dtype))

    num_values = int(num_values)
    maximum = num_values - 1
    dtype = np.dtype(dtype)

    if np.min_scalar_type(maximum) > dtype:
      raise ValueError(_DTYPE_OVERFLOW.format(dtype, num_values))

    super(DiscreteArray, self).__init__(
        shape=(),
        dtype=dtype,
        minimum=0,
        maximum=maximum,
        name=name)
    self._num_values = num_values

  @property
  def num_values(self):
    """Returns the number of items."""
    return self._num_values

  def __repr__(self):
    return self._REPR_TEMPLATE.format(self=self)  # pytype: disable=duplicate-keyword-argument

  def __reduce__(self):
    return DiscreteArray, (self._num_values, self._dtype, self._name)


_VALID_STRING_TYPES = (str, bytes)
_INVALID_STRING_TYPE = (
    'Expected `string_type` to be one of: {}, got: {{!r}}.'
    .format(_VALID_STRING_TYPES))
_INVALID_ELEMENT_TYPE = (
    'Expected all elements to be of type: %s. Got value: %r of type: %s.')


class StringArray(Array):
  """Represents an array of variable-length Python strings."""
  __slots__ = ('_string_type',)

  _REPR_TEMPLATE = (
      '{self.__class__.__name__}(shape={self.shape}, '
      'string_type={self.string_type}, name={self.name})')

  def __init__(self, shape, string_type=str, name=None):
    """Initializes a new `StringArray` spec.

    Args:
      shape: An iterable specifying the array shape.
      string_type: The native Python string type for each element; either
        unicode or ASCII. Defaults to unicode.
      name: Optional string containing a semantic name for the corresponding
        array. Defaults to `None`.
    """
    if string_type not in _VALID_STRING_TYPES:
      raise ValueError(_INVALID_STRING_TYPE.format(string_type))
    self._string_type = string_type
    super(StringArray, self).__init__(shape=shape, dtype=object, name=name)

  @property
  def string_type(self):
    """Returns the Python string type for each element."""
    return self._string_type

  def validate(self, value):
    """Checks if value conforms to this spec.

    Args:
      value: a numpy array or value convertible to one via `np.asarray`.

    Returns:
      value, converted if necessary to a numpy array.

    Raises:
      ValueError: if value doesn't conform to this spec.
    """
    value = np.asarray(value, dtype=object)
    if value.shape != self.shape:
      self._fail_validation(_INVALID_SHAPE, self.shape, value.shape)
    for item in value.flat:
      if not isinstance(item, self.string_type):
        self._fail_validation(
            _INVALID_ELEMENT_TYPE, self.string_type, item, type(item))
    return value

  def generate_value(self):
    """Generate a test value which conforms to this spec."""
    empty_string = self.string_type()  # pylint: disable=not-callable
    return np.full(shape=self.shape, dtype=self.dtype, fill_value=empty_string)

  def __repr__(self):
    return self._REPR_TEMPLATE.format(self=self)  # pytype: disable=duplicate-keyword-argument

  def __reduce__(self):
    return type(self), (self.shape, self.string_type, self.name)
