# coding=utf-8
# Copyright 2020 The TF-Agents 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
#
#     https://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.

"""Utilities related to TensorSpec class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
from typing import Union

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

from tf_agents.specs import array_spec
from tf_agents.typing import types

from google.protobuf import text_format
# pylint:disable=g-direct-tensorflow-import
from tensorflow.core.protobuf import struct_pb2  # TF internal
from tensorflow.python.framework import tensor_spec as ts  # TF internal
from tensorflow.python.saved_model import nested_structure_coder  # TF internal
# pylint:enable=g-direct-tensorflow-import

tfd = tfp.distributions

TensorSpec = tf.TensorSpec
BoundedTensorSpec = ts.BoundedTensorSpec


def is_bounded(spec):
  if isinstance(spec, (array_spec.BoundedArraySpec, BoundedTensorSpec)):
    return True
  elif hasattr(spec, "minimum") and hasattr(spec, "maximum"):
    return hasattr(spec, "dtype") and hasattr(spec, "shape")


def is_discrete(spec):
  if isinstance(spec, TensorSpec):
    return spec.dtype.is_integer
  else:
    return array_spec.is_discrete(spec)


def is_continuous(spec):
  if isinstance(spec, TensorSpec):
    return spec.dtype.is_floating
  else:
    return array_spec.is_continuous(spec)


def from_spec(spec):
  """Maps the given spec into corresponding TensorSpecs keeping bounds."""

  def _convert_to_tensor_spec(s):
    # Need to check bounded first as non bounded specs are base class.
    if isinstance(s, tf.TypeSpec):
      return s
    if is_bounded(s):
      return BoundedTensorSpec.from_spec(s)
    elif isinstance(s, array_spec.ArraySpec):
      return TensorSpec.from_spec(s)
    else:
      raise ValueError(
          "No known conversion from type `%s` to a TensorSpec.  Saw:\n  %s"
          % (type(s), s))

  return tf.nest.map_structure(_convert_to_tensor_spec, spec)


def to_array_spec(
    tensor_spec: Union[types.NestedArraySpec, types.NestedTensorSpec]
) -> types.NestedArraySpec:
  """Converts TensorSpec into ArraySpec."""
  def _convert(s):
    if isinstance(s, array_spec.ArraySpec):
      return s

    if hasattr(s, "minimum") and hasattr(s, "maximum"):
      return array_spec.BoundedArraySpec(
          s.shape.as_list(),
          s.dtype.as_numpy_dtype,
          minimum=s.minimum,
          maximum=s.maximum,
          name=s.name)
    else:
      return array_spec.ArraySpec(s.shape.as_list(),
                                  s.dtype.as_numpy_dtype,
                                  s.name)

  return tf.nest.map_structure(_convert, tensor_spec)


def to_nest_array_spec(
    nest_array_spec: Union[types.NestedArraySpec, types.NestedTensorSpec]
) -> types.NestedArraySpec:
  """(Deprecated) Alias for `to_array_spec`."""
  return to_array_spec(nest_array_spec)


def to_placeholder(spec, outer_dims=()):
  """Creates a placeholder from TensorSpec.

  Args:
    spec: instance of TensorSpec
    outer_dims: optional leading dimensions of the placeholder.

  Returns:
    An instance of tf.placeholder.
  """
  ph_shape = list(outer_dims) + spec.shape.as_list()
  return tf.compat.v1.placeholder(spec.dtype, ph_shape, spec.name)


def to_placeholder_with_default(default, spec, outer_dims=()):
  """Creates a placeholder from TensorSpec.

  Args:
    default: A constant value of output type dtype.
    spec: Instance of TensorSpec
    outer_dims: Optional leading dimensions of the placeholder.

  Returns:
    An instance of tf.placeholder.
  """
  ph_shape = list(outer_dims) + spec.shape.as_list()
  return tf.compat.v1.placeholder_with_default(default, ph_shape, spec.name)


def to_nest_placeholder(nested_tensor_specs,
                        default=None,
                        name_scope="",
                        outer_dims=()):
  """Converts a nest of TensorSpecs to a nest of matching placeholders.

  Args:
    nested_tensor_specs: A nest of tensor specs.
    default: Optional constant value to set as a default for the placeholder.
    name_scope: String name for the scope to create the placeholders in.
    outer_dims: Optional leading dimensions for the placeholder.

  Returns:
    A nest of placeholders matching the given tensor spec.

  Raises:
    ValueError: If a default is provided outside of the allowed types, or if
      default is a np.array that does not match the spec shape.
  """
  if default is None:
    to_ph = lambda spec: to_placeholder(spec, outer_dims=outer_dims)
  else:
    if not isinstance(default, (int, float, np.ndarray)):
      raise ValueError("to_nest_placeholder default value must be an int, "
                       "float, or np.ndarray")

    def to_ph(spec):
      shape = list(outer_dims) + spec.shape.as_list()
      if isinstance(default, np.ndarray) and list(default.shape) != shape:
        raise ValueError("Shape mismatch between default value and spec. "
                         "Got {}, expected {}".format(default.shape, shape))
      const = tf.constant(default, shape=shape, dtype=spec.dtype)
      return to_placeholder_with_default(const, spec, outer_dims=outer_dims)

  with tf.name_scope(name_scope):
    return tf.nest.map_structure(to_ph, nested_tensor_specs)


def _random_uniform_int(shape, outer_dims, minval, maxval, dtype, seed=None):
  """Iterates over n-d tensor minval, maxval limits to sample uniformly."""
  # maxval in BoundedTensorSpec is bound inclusive.
  # tf.random_uniform is upper bound exclusive, +1 to fix the sampling
  # behavior.
  # However +1 could cause overflow, in such cases we use the original maxval.
  maxval = np.broadcast_to(maxval, minval.shape).astype(dtype.as_numpy_dtype)
  minval = np.broadcast_to(minval, maxval.shape).astype(dtype.as_numpy_dtype)

  sampling_maxval = maxval
  if dtype.is_integer:
    sampling_maxval = np.where(maxval < dtype.max, maxval + 1, maxval)

  if not np.all(shape[-len(minval.shape):] == minval.shape):
    raise ValueError(
        "%s == shape[-%d:] != minval.shape == %s.  shape == %s." %
        (shape[len(minval.shape):], len(minval.shape), minval.shape, shape))

  # Example:
  #  minval = [1.0, 2.0]
  #  shape = [3, 2]
  #  outer_dims = [5]
  # Sampling becomes:
  #  sample [5, 3] for minval 1.0
  #  sample [5, 3] for minval 2.0
  #  stack on innermost axis to get [5, 3, 2]
  #  reshape to get [5, 3, 2]
  samples = []
  shape = tf.convert_to_tensor(shape, dtype=tf.int32)
  sample_shape = tf.concat((outer_dims, shape[:-len(minval.shape)]), axis=0)
  full_shape = tf.concat((outer_dims, shape), axis=0)
  for (single_min, single_max) in zip(minval.flat, sampling_maxval.flat):
    samples.append(
        tf.random.uniform(
            shape=sample_shape,
            minval=single_min,
            maxval=single_max,
            dtype=dtype,
            seed=seed))
  samples = tf.stack(samples, axis=-1)
  samples = tf.reshape(samples, full_shape)
  return samples


def sample_bounded_spec(spec, seed=None, outer_dims=None):
  """Samples uniformily the given bounded spec.

  Args:
    spec: A BoundedSpec to sample.
    seed: A seed used for sampling ops
    outer_dims: An optional `Tensor` specifying outer dimensions to add to the
      spec shape before sampling.

  Returns:
    A Tensor sample of the requested spec.
  """
  minval = spec.minimum
  maxval = spec.maximum
  dtype = tf.as_dtype(spec.dtype)

  # To sample uint8 we will use int32 and cast later. This is needed for two
  # reasons:
  #  - tf.random_uniform does not currently support uint8
  #  - if you want to sample [0, 255] range, there's no way to do this since
  #    tf.random_uniform has exclusive upper bound and 255 + 1 would overflow.
  is_uint8 = dtype == tf.uint8
  sampling_dtype = tf.int32 if is_uint8 else dtype

  if dtype in [tf.float64, tf.float32]:
    # Avoid under/over-flow as random_uniform can't sample over the full range
    # for these types.
    minval = np.maximum(dtype.min / 8, minval)
    maxval = np.minimum(dtype.max / 8, maxval)

  if outer_dims is None:
    outer_dims = tf.constant([], dtype=tf.int32)
  else:
    outer_dims = tf.convert_to_tensor(outer_dims, dtype=tf.int32)

  def _unique_vals(vals):
    if vals.size > 0:
      if vals.ndim > 0:
        return np.all(vals == vals[0])
    return True

  if (minval.ndim != 0 or
      maxval.ndim != 0) and not (_unique_vals(minval) and _unique_vals(maxval)):
    # tf.random_uniform can only handle minval/maxval 0-d tensors.
    res = _random_uniform_int(
        shape=spec.shape,
        outer_dims=outer_dims,
        minval=minval,
        maxval=maxval,
        dtype=sampling_dtype,
        seed=seed)
  else:
    minval = minval.item(0) if minval.ndim != 0 else minval
    maxval = maxval.item(0) if maxval.ndim != 0 else maxval
    # BoundedTensorSpec are bounds inclusive.
    # tf.random_uniform is upper bound exclusive, +1 to fix the sampling
    # behavior.
    # However +1 will cause overflow, in such cases we use the original maxval.
    if sampling_dtype.is_integer and maxval < sampling_dtype.max:
      maxval = maxval + 1

    shape = tf.convert_to_tensor(spec.shape, dtype=tf.int32)
    full_shape = tf.concat((outer_dims, shape), axis=0)
    res = tf.random.uniform(
        full_shape,
        minval=minval,
        maxval=maxval,
        dtype=sampling_dtype,
        seed=seed)

  if is_uint8:
    res = tf.cast(res, dtype=dtype)

  return res


def sample_spec_nest(structure,
                     seed=None,
                     outer_dims=(),
                     minimum=None,
                     maximum=None):
  """Samples the given nest of specs.

  Args:
    structure: A nest of `TensorSpec`.
    seed: A seed used for sampling ops
    outer_dims: An optional `Tensor` specifying outer dimensions to add to the
      spec shape before sampling.
    minimum: An optional numeric value. If set, numeric specs within the nest
      (both bounded and unbounded) will be restricted to this minimum.
    maximum: Similar to the above but with maximums.

  Returns:
    A nest of sampled values following the ArraySpec definition.

  Raises:
    TypeError: If `spec` is an unknown type.
    NotImplementedError: If `outer_dims` is not statically known but nest
      contains a `SparseTensorSpec`.
  """
  seed_stream = tfp.util.SeedStream(seed=seed, salt="sample_spec_nest")

  def sample_fn(spec):
    """Return a composite tensor sample given `spec`.

    Args:
      spec: A TensorSpec, SparseTensorSpec, etc.

    Returns:
      A tensor or SparseTensor.

    Raises:
      NotImplementedError: If `outer_dims` is not statically known and a
        SparseTensor is requested.
    """
    if isinstance(spec, tf.SparseTensorSpec):
      outer_shape = tf.get_static_value(outer_dims)
      if outer_dims is not None and outer_shape is None:
        raise NotImplementedError(
            "outer_dims must be statically known, got: {}".format(outer_dims))
      shape = tf.TensorShape(outer_shape or []).concatenate(spec.shape)

      if shape.num_elements() == 0 or tf.compat.dimension_value(shape[0]) == 0:
        return tf.SparseTensor(
            indices=tf.zeros([0, shape.rank], dtype=tf.int64),
            values=tf.zeros([0], dtype=spec.dtype),
            dense_shape=shape)

      indices_spec = BoundedTensorSpec(
          dtype=tf.int64,
          shape=[7, shape.rank],
          minimum=[0] * shape.rank,
          maximum=[x - 1 for x in shape.as_list()])
      values_dtype = tf.int32 if spec.dtype == tf.string else spec.dtype
      values_spec = BoundedTensorSpec(
          dtype=values_dtype,
          shape=[7],
          minimum=0,
          maximum=shape.as_list()[-1] - 1)
      values_sample = sample_bounded_spec(values_spec, seed=seed_stream())
      if spec.dtype == tf.string:
        values_sample = tf.as_string(values_sample)
      return tf.sparse.reorder(
          tf.SparseTensor(
              indices=sample_bounded_spec(indices_spec, seed=seed_stream()),
              values=values_sample,
              dense_shape=shape))
    elif isinstance(spec, (TensorSpec, BoundedTensorSpec)):
      if spec.dtype == tf.string:
        sample_spec = BoundedTensorSpec(
            spec.shape, tf.int32, minimum=0, maximum=10)
        return tf.as_string(
            sample_bounded_spec(
                sample_spec, outer_dims=outer_dims, seed=seed_stream()))
      elif spec.dtype == tf.bool:
        sample_spec = BoundedTensorSpec(
            spec.shape, tf.int32, minimum=0, maximum=1)
        return tf.cast(sample_bounded_spec(
            sample_spec, outer_dims=outer_dims, seed=seed_stream()), tf.bool)
      else:
        bounded_spec = BoundedTensorSpec.from_spec(spec)

        spec_max = bounded_spec.maximum
        if maximum is not None:
          spec_max = min(maximum, spec_max)

        spec_min = bounded_spec.minimum
        if minimum is not None:
          spec_min = max(minimum, spec_min)
        bounded_spec = BoundedTensorSpec(
            shape=bounded_spec.shape,
            dtype=bounded_spec.dtype,
            minimum=spec_min,
            maximum=spec_max)
        return sample_bounded_spec(
            bounded_spec,
            outer_dims=outer_dims,
            seed=seed_stream(),
        )
    else:
      raise TypeError("Spec type not supported: '{}'".format(spec))

  return tf.nest.map_structure(sample_fn, structure)


def zero_spec_nest(specs, outer_dims=None):
  """Create zero tensors for a given spec.

  Args:
    specs: A nest of `TensorSpec`.
    outer_dims: An optional list of constants or `Tensor` specifying outer
      dimensions to add to the spec shape before sampling.

  Returns:
    A nest of zero tensors matching `specs`, with the optional outer
    dimensions added.

  Raises:
    TypeError: If `specs` is an unknown type.
    NotImplementedError: If `specs` contains non-dense tensor specs.
  """

  def make_zero(spec):
    if not isinstance(spec, TensorSpec):
      raise NotImplementedError("Spec type not supported: '{}'".format(spec))
    if outer_dims is None:
      shape = spec.shape
    else:
      spec_shape = tf.convert_to_tensor(value=spec.shape, dtype=tf.int32)
      shape = tf.concat((outer_dims, spec_shape), axis=0)
    return tf.zeros(shape, spec.dtype)

  if specs:
    if outer_dims is None:
      outer_dims = tf.constant([], dtype=tf.int32)
    else:
      outer_dims = tf.convert_to_tensor(outer_dims, dtype=tf.int32)

  return tf.nest.map_structure(make_zero, specs)


def add_outer_dims_nest(specs, outer_dims):
  """Adds outer dimensions to the shape of input specs.

  Args:
    specs: Nested list/tuple/dict of TensorSpecs/ArraySpecs, describing the
      shape of tensors.
    outer_dims: a list or tuple, representing the outer shape to be added to the
      TensorSpecs in specs.

  Returns:
    Nested TensorSpecs with outer dimensions added to the shape of input specs.

  Raises:
    ValueError: if any outer_dims is neither a list nor tuple.
  """
  if not isinstance(outer_dims, (tuple, list)):
    raise ValueError("outer_dims must be a tuple or list of dimensions")

  def add_outer_dims(spec):
    # TODO(b/187478998) Use spec.name when tf.SparseTensorSpec supports it.
    name = getattr(spec, "name", None)
    shape = outer_dims + spec.shape
    if hasattr(spec, "minimum") and hasattr(spec, "maximum"):
      return BoundedTensorSpec(shape, spec.dtype, spec.minimum, spec.maximum,
                               name)
    elif isinstance(specs, tf.SparseTensorSpec):
      # TODO(b/187478998) Add name when tf.SparseTensorSpec supports it.
      return tf.SparseTensorSpec(shape, spec.dtype)
    return TensorSpec(shape, spec.dtype, name=name)

  return tf.nest.map_structure(add_outer_dims, specs)


def add_outer_dim(specs, dim=None):
  """Adds an outer dimension to the shape of input specs.

  Args:
    specs: Nested list/tuple/dict of TensorSpecs/ArraySpecs, describing the
      shape of tensors.
    dim: Int, representing the outer dimension to be added to the TensorSpecs in
      specs.

  Returns:
    Nested TensorSpecs with outer dimensions added to the shape of input specs.

  """
  return add_outer_dims_nest(specs, outer_dims=(dim,))


def with_dtype(specs, dtype):
  """Updates dtypes of all specs in the input spec.

  Args:
    specs: Nested list/tuple/dict of TensorSpecs/ArraySpecs, describing the
      shape of tensors.
    dtype: dtype to update the specs to.

  Returns:
    Nested TensorSpecs with the udpated dtype.
  """

  def update_dtype(spec):
    if hasattr(spec, "minimum") and hasattr(spec, "maximum"):
      return BoundedTensorSpec(spec.shape, dtype, spec.minimum, spec.maximum,
                               spec.name)
    return TensorSpec(spec.shape, dtype, name=spec.name)

  return tf.nest.map_structure(update_dtype, specs)


def remove_outer_dims_nest(specs, num_outer_dims):
  """Removes the specified number of outer dimensions from the input spec nest.

  Args:
    specs: Nested list/tuple/dict of TensorSpecs/ArraySpecs, describing the
      shape of tensors.
    num_outer_dims: (int) Number of outer dimensions to remove.

  Returns:
    Nested TensorSpecs with outer dimensions removed from the input specs.

  Raises:
    Value error if a spec in the nest has shape rank less than `num_outer_dims`.
  """

  def remove_outer_dims(spec):
    """Removes the num_outer_dims of a tensor spec."""
    # TODO(b/187478998) Use spec.name when tf.SparseTensorSpec supports it.
    name = getattr(spec, "name", None)
    if len(spec.shape) < num_outer_dims:
      raise ValueError("The shape of spec {} has rank lower than the specified "
                       "num_outer_dims {}".format(spec, num_outer_dims))
    shape = list(spec.shape)[num_outer_dims:]
    if hasattr(spec, "minimum") and hasattr(spec, "maximum"):
      if isinstance(spec.minimum,
                    (tuple, list)) and len(spec.minimum) == len(spec.shape):
        minimum = spec.minimum[num_outer_dims:]
      else:
        minimum = spec.minimum
      if isinstance(spec.maximum,
                    (tuple, list)) and len(spec.maximum) == len(spec.shape):
        maximum = spec.maximum[num_outer_dims:]
      else:
        maximum = spec.maximum

      return BoundedTensorSpec(shape, spec.dtype, minimum, maximum, name)
    elif isinstance(spec, tf.SparseTensorSpec):
      # TODO(b/187478998) Add name when tf.SparseTensorSpec supports it.
      return tf.SparseTensorSpec(shape, spec.dtype)
    return TensorSpec(shape, spec.dtype, name=name)

  return tf.nest.map_structure(remove_outer_dims, specs)


def to_proto(spec):
  """Encodes a nested spec into a struct_pb2.StructuredValue proto.

  Args:
    spec: Nested list/tuple or dict of TensorSpecs, describing the
      shape of the non-batched Tensors.
  Returns:
    A `struct_pb2.StructuredValue` proto.
  """
  # Make sure spec is a tensor_spec.
  spec = from_spec(spec)
  return nested_structure_coder.encode_structure(spec)


def from_proto(spec_proto):
  """Decodes a struct_pb2.StructuredValue proto into a nested spec."""
  return nested_structure_coder.decode_proto(spec_proto)


def from_packed_proto(spec_packed_proto):
  """Decodes a packed Any proto containing the structured value for the spec."""
  spec_proto = struct_pb2.StructuredValue()
  spec_packed_proto.Unpack(spec_proto)
  return from_proto(spec_proto)


def to_pbtxt_file(output_path, spec):
  """Saves a spec encoded as a struct_pb2.StructuredValue in a pbtxt file."""
  spec_proto = to_proto(spec)
  dir_path = os.path.split(output_path)[0]
  tf.io.gfile.makedirs(dir_path)
  with tf.io.gfile.GFile(output_path, "wb") as f:
    f.write(text_format.MessageToString(spec_proto))


def from_pbtxt_file(spec_path):
  """Loads a spec encoded as a struct_pb2.StructuredValue from a pbtxt file."""
  spec_proto = struct_pb2.StructuredValue()
  with tf.io.gfile.GFile(spec_path, "rb") as f:
    text_format.MergeLines(f, spec_proto)
  return from_proto(spec_proto)
