# 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 distributions."""

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

import inspect
from typing import Any, Mapping, Type, Text

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

from tf_agents.distributions import tanh_bijector_stable
from tf_agents.specs import tensor_spec
from tf_agents.typing import types
from tf_agents.utils import common
from tf_agents.utils import nest_utils


def scale_distribution_to_spec(distribution, spec):
  """Scales the given distribution to the bounds of the given spec."""
  return SquashToSpecNormal(distribution, spec)


class SquashToSpecNormal(tfp.distributions.AutoCompositeTensorDistribution):
  """Scales an input normalized action distribution to match spec bounds.

  Unlike the normal distribution computed when NormalProjectionNetwork
  is called with scale_distribution=False, which merely squashes the mean
  of the distribution to within the action spec, this distribution scales the
  output distribution to ensure that the output action fits within the spec.

  This distribution also maintains the input normal distribution, and uses this
  distribution to compute the KL-divergence between two SquashToSpecNormal
  distributions provided that they were scaled by the same action spec.
  This is possible as KL divergence is invariant when both distributions are
  transformed using the same invertible function.

  Formally, let a be the action magnitude and b be the action mean. The
  squashing operation performs the following change of variables to the
  input distribution X:

  Y = a * tanh(X) + b

  Note that this is a change of variables as the function is invertible, with:

  X = tan((Y - b) / a), where Y in (b - a, b + a)
  """

  def __init__(self,
               distribution,
               spec,
               validate_args=False,
               name="SquashToSpecNormal"):
    """Constructs a SquashToSpecNormal distribution.

    Args:
      distribution: input normal distribution with normalized mean and std dev
      spec: bounded action spec from which to compute action ranges
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      name: Python `str` name prefixed to Ops created by this class.
    """
    parameters = dict(locals())
    if not isinstance(
        distribution,
        (tfp.distributions.Normal, tfp.distributions.MultivariateNormalDiag)):
      raise ValueError("Input distribution must be a normal distribution, "
                       "got {} instead".format(distribution))
    self.action_means, self.action_magnitudes = common.spec_means_and_magnitudes(
        spec)
    self._distribution = distribution
    self._spec = spec

    bijectors = [
        tfp.bijectors.Shift(self.action_means)(
            tfp.bijectors.Scale(self.action_magnitudes)),
        tanh_bijector_stable.Tanh()
    ]
    bijector_chain = tfp.bijectors.Chain(bijectors)
    self._squashed_distribution = tfp.distributions.TransformedDistribution(
        distribution=distribution, bijector=bijector_chain)
    super().__init__(
        dtype=distribution.dtype,
        reparameterization_type=distribution.reparameterization_type,
        validate_args=validate_args,
        allow_nan_stats=distribution.allow_nan_stats,
        parameters=parameters,
        # We let TransformedDistribution access _graph_parents since this class
        # is more like a baseclass than derived.
        graph_parents=(
            distribution._graph_parents +  # pylint: disable=protected-access
            bijector_chain.graph_parents),
        name=name)

  @classmethod
  def _parameter_properties(cls, dtype, num_classes=None):
    return dict(
        distribution=tfp.util.BatchedComponentProperties())

  @property
  def input_distribution(self):
    """The raw action distribution."""
    return self._distribution

  def kl_divergence(self, other, name="kl_divergence"):
    """Computes the KL Divergence between two SquashToSpecNormal distributions."""
    if not isinstance(other, SquashToSpecNormal):
      raise ValueError("other distribution should be of type "
                       "SquashToSpecNormal, got {}".format(other))
    if (np.any(self.action_means != other.action_means) or
        np.any(self.action_magnitudes != other.action_magnitudes)):
      raise ValueError("Other distribution does not have same action mean "
                       "and magnitude. This mean {}, this magnitude {}, "
                       "other mean {}, other magnitude {}.".format(
                           self.action_means, self.action_magnitudes,
                           other.action_means, other.action_magnitudes))
    return self.input_distribution.kl_divergence(other.input_distribution, name)

  def sample(self, sample_shape=(), seed=None, name="sample"):
    """Generates samples from the wrapped TransformedDistribution."""
    return self._squashed_distribution.sample(sample_shape, seed, name)

  def log_prob(self, value, name="log_prob"):
    """Computes log probability from the wrapped TransformedDistribution."""
    return self._squashed_distribution.log_prob(value, name)

  def prob(self, value, name="prob"):
    """Computes probability from the wrapped TransformedDistribution."""
    return self._squashed_distribution.prob(value, name)

  def stddev(self, name="stddev"):
    """Compute stddev of the SquashToSpecNormal distribution."""
    stddev = self.action_magnitudes * tf.tanh(self.input_distribution.stddev())
    return stddev

  def mode(self, name="mode"):
    """Compute mean of the SquashToSpecNormal distribution."""
    mean = self.action_magnitudes * tf.tanh(self.input_distribution.mode()) + \
        self.action_means
    return mean

  def mean(self, name="mean", **kwargs):
    """Compute mean of the SquashToSpecNormal distribution."""
    return self.mode(name)

  def event_shape_tensor(self, name="event_shape_tensor"):
    """Compute event shape tensor of the SquashToSpecNormal distribution."""
    return self._squashed_distribution.event_shape_tensor(name)

  def batch_shape_tensor(self, name="batch_shape_tensor"):
    """Compute event shape tensor of the SquashToSpecNormal distribution."""
    return self._squashed_distribution.batch_shape_tensor(name)


class Params(object):
  """The (recursive) parameters of objects exposing the `parameters` property.

  This includes TFP `Distribution`, `Bijector`, and TF `LinearOperator`.

  `Params` objects are created with
  `tf_agents.distributions.utils.get_parameters`;
  `Params` can be converted back to original objects via
  `tf_agents.distributions.utils.make_from_parameters`.

  In-place edits of fields are allowed, and will not modify the original
  objects (with the exception of, e.g., reference objects like `tf.Variable`
  being modified in-place).

  The components of a `Params` object are: `type_` and `params`.

  - `type_` is the type of object.
  - `params` is a `dict` of the (non-default) non-tensor arguments passed to the
    object's `__init__`; and includes nests of Python objects, as well as other
    `Params` values representing "Param-representable" objects passed to init.

  A non-trivial example:

  ```python
  scale_matrix = tf.Variable([[1.0, 2.0], [-1.0, 0.0]])
  d = tfp.distributions.MultivariateNormalDiag(
      loc=[1.0, 1.0], scale_diag=[2.0, 3.0], validate_args=True)
  b = tfp.bijectors.ScaleMatvecLinearOperator(
      scale=tf.linalg.LinearOperatorFullMatrix(matrix=scale_matrix),
      adjoint=True)
  b_d = b(d)
  p = utils.get_parameters(b_d)
  ```

  Then `p` is:

  ```python
  Params(
      tfp.distributions.TransformedDistribution,
      params={
          "bijector": Params(
              tfp.bijectors.ScaleMatvecLinearOperator,
              params={"adjoint": True,
                      "scale": Params(
                          tf.linalg.LinearOperatorFullMatrix,
                          params={"matrix": scale_matrix})}),
          "distribution": Params(
              tfp.distributions.MultivariateNormalDiag,
              params={"validate_args": True,
                      "scale_diag": [2.0, 3.0],
                      "loc": [1.0, 1.0]})})
  ```

  This structure can be manipulated and/or converted back to a `Distribution`
  instance via `make_from_parameters`:

  ```python
  p.params["distribution"].params["loc"] = [0.0, 0.0]

  # The distribution `new_b_d` will be a MVN centered on `(0, 0)` passed through
  # the `ScaleMatvecLinearOperator` bijector.
  new_b_d = utils.make_from_parameters(p)
  ```
  """
  type_: Type[Any]  # Any class that has a .parameters.
  params: Mapping[Text, Any]

  def __str__(self):
    return "<Params: type={}, params={}>".format(self.type_, self.params)

  def __repr__(self):
    return str(self)

  def __eq__(self, other):
    return (isinstance(self, type(other))
            and self.type_ == other.type_
            and self.params == other.params)

  def __init__(self, type_, params):
    self.type_ = type_
    self.params = params


def get_parameters(value: Any) -> Params:
  """Creates a recursive `Params` object from `value`.

  The `Params` object can be converted back to an object of type `type(value)`
  with `make_from_parameters`.  For more details, see the docstring of
  `Params`.

  Args:
    value: Typically a user provides `tfp.Distribution`, `tfp.Bijector`, or
      `tf.linalg.LinearOperator`, but this can be any Python object.

  Returns:
    An instance of `Params`.

  Raises:
    TypeError: If `value.parameters` exists, is not `None`, but but is also not
       a `Mapping` (e.g. a `dict`).
  """
  parameters = getattr(value, "parameters", None)
  if not isinstance(parameters, Mapping):
    raise TypeError(
        "value.parameters is not available or is not a dict; "
        "value: {}; parameters: {}".format(value, parameters))
  type_ = type(value)
  params = {}

  def process_parameter(p):
    if getattr(p, "parameters", None) is not None:
      return get_parameters(p)
    else:
      return p

  if getattr(value, "parameters"):
    default_values = inspect.signature(type_.__init__).parameters.items()
    default_values = {
        k: v.default
        for (k, v) in default_values
        if v.default is not inspect.Parameter.empty
    }
    params = {
        k: tf.nest.map_structure(process_parameter, v)
        for k, v in value.parameters.items()
        if v is not default_values.get(k, None)
    }
    return Params(type(value), params)

  return value


def make_from_parameters(value: Params) -> Any:
  """Creates an instance of type `value.type_` with the parameters in `value`.

  For more details, see the docstrings for `get_parameters` and `Params`.

  This function may raise strange errors if `value` is a `Params` created from
  a badly constructed object (one which does not set `self._parameters`
  properly).  For example:

  ```python
  class MyBadlyConstructedDistribution(tfp.distributions.Categorical):
    def __init__(self, extra_arg, **kwargs):
      super().__init__(**kwargs)
      self._extra_arg = extra_arg

    ...
  ```

  To fix this, make sure `self._parameters` are properly set:

  ```python
  class MyProperlyConstructedDistribution(tfp.distributions.Categorical):
    def __init__(self, extra_arg, **kwargs):
      super().__init__(**kwargs)
      # Ensure all arguments to `__init__` are in `self._parameters`.
      self._parameters = dict(extra_arg=extra_arg, **kwargs)
      self._extra_arg = extra_arg

    ...
  ```

  Args:
    value: A `Params` object; the output of `get_parameters` (or a
      modified version thereof).

  Returns:
    An instance of `value.type_`.

  Raises:
    Exception: If `value` is a `Params` object and the initializer of
      `value.type_` does not recognize accept the args structure given in
      `value.params`.  This can happen if, e.g., `value.type_.__init__` does not
      properly set `self._parameters` or `self.parameters` to match the
      arguments it expects.
  """
  def make_from_params_or_identity(v_):
    return make_from_parameters(v_) if isinstance(v_, Params) else v_

  params = {
      k: tf.nest.map_structure(make_from_params_or_identity, v)
      for k, v in value.params.items()
  }
  return value.type_(**params)


def parameters_to_dict(
    value: Params,
    tensors_only: bool = False) -> Mapping[Text, Any]:
  """Converts `value` to a nested `dict` (excluding all `type_` info).

  Sub-dicts represent `Params` objects; keys represent flattened nest structures
  in `value.params`.

  Example:

  ```python
  scale_matrix = tf.Variable([[1.0, 2.0], [-1.0, 0.0]])
  d = tfp.distributions.MultivariateNormalDiag(
      loc=[1.0, 1.0], scale_diag=[2.0, 3.0], validate_args=True)
  b = tfp.bijectors.ScaleMatvecLinearOperator(
      scale=tf.linalg.LinearOperatorFullMatrix(matrix=scale_matrix),
      adjoint=True)
  b_d = b(d)
  p = utils.get_parameters(b_d)
  params_dict = utils.parameters_to_dict(p, tensors_only)
  ```

  results in the nested dictionary, if `tensors_only=False`:

  ```python
  {
    "bijector": {"adjoint": True,
                 "scale": {"matrix": scale_matrix}},
    "distribution": {"validate_args": True,
                     # These are deeply nested because we passed lists
                     # intead of numpy arrays for `loc` and `scale_diag`.
                     "scale_diag:0": 2.0,
                     "scale_diag:1": 3.0,
                     "loc:0": 1.0,
                     "loc:1": 1.0}
  }
  ```

  results if `tensors_only=True`:

  ```python
  {
    "bijector": {"scale": {"matrix": scale_matrix}},
    "distribution": {},
  }
  ```

  The dictionary may then be modified or updated (e.g., in place), and converted
  back to a `Params` object using `merge_to_parameters_from_dict`.

  Args:
    value: The (possibly recursively defined) `Params`.
    tensors_only: Whether to include all parameters or only
      tensors and TypeSpecs.  If `True`, then only
      `tf.Tensor` and `tf.TypeSpec` leaf parameters are emitted.

  Returns:
    A `dict` mapping `value.params` to flattened key/value pairs.  Any
    sub-`Params` objects become nested dicts.
  """
  output_entries = {}

  # All tensors and TensorSpecs will be included for purposes of this test.
  is_tensor_or_spec = (
      lambda p: tf.is_tensor(p) or isinstance(p, tf.TypeSpec))

  def convert(key, p):
    if isinstance(p, Params):
      output_entries[key] = parameters_to_dict(p, tensors_only)
    elif not tensors_only or is_tensor_or_spec(p):
      output_entries[key] = p

  for k, v in value.params.items():
    if tf.nest.is_nested(v):
      flattened_params = nest_utils.flatten_with_joined_paths(v)
      for (param_k, param_v) in flattened_params:
        key = "{}:{}".format(k, param_k)
        convert(key, param_v)
    else:
      convert(k, v)

  return output_entries


def merge_to_parameters_from_dict(
    value: Params, params_dict: Mapping[Text, Any]) -> Params:
  """Merges dict matching data of `parameters_to_dict(value)` to a new `Params`.

  For more details, see the example below and the documentation of
  `parameters_to_dict`.

  Example:

  ```python
  scale_matrix = tf.Variable([[1.0, 2.0], [-1.0, 0.0]])
  d = tfp.distributions.MultivariateNormalDiag(
      loc=[1.0, 1.0], scale_diag=[2.0, 3.0], validate_args=True)
  b = tfp.bijectors.ScaleMatvecLinearOperator(
      scale=tf.linalg.LinearOperatorFullMatrix(matrix=scale_matrix),
      adjoint=True)
  b_d = b(d)
  p = utils.get_parameters(b_d)

  params_dict = utils.parameters_to_dict(p)
  params_dict["bijector"]["scale"]["matrix"] = new_scale_matrix

  new_params = utils.merge_to_parameters_from_dict(
    p, params_dict)

  # new_d is a `ScaleMatvecLinearOperator()(MultivariateNormalDiag)` with
  # a new scale matrix.
  new_d = utils.make_from_parameters(new_params)
  ```

  Args:
    value: A `Params` from which `params_dict` was derived.
    params_dict: A nested `dict` created by e.g. calling
      `parameters_to_dict(value)` and  modifying it to modify parameters.
      **NOTE** If any keys in the dict are missing, the "default" value in
      `value` is used instead.

  Returns:
    A new `Params` object which can then be turned into e.g. a
    `tfp.Distribution` via `make_from_parameters`.

  Raises:
    ValueError: If `params_dict` has keys missing from `value.params`.
    KeyError: If a subdict entry is missing for a nested value in
      `value.params`.
  """
  new_params = {}
  if params_dict is None:
    params_dict = {}

  processed_params = set()
  for k, v in value.params.items():
    # pylint: disable=cell-var-from-loop
    visited = set()
    converted = set()

    def convert(params_k, p):
      if params_k is not None:
        params_key = "{}:{}".format(k, params_k)
        visited.add(params_key)
        params_dict_value = params_dict.get(params_key, None)
        if params_dict_value is not None:
          converted.add(params_key)
      else:
        params_key = k
        params_dict_value = params_dict.get(k, None)
      processed_params.add(params_key)
      if isinstance(p, Params):
        return merge_to_parameters_from_dict(p, params_dict_value)
      else:
        return params_dict_value if params_dict_value is not None else p
    # pylint: enable=cell-var-from-loop

    if tf.nest.is_nested(v):
      new_params[k] = nest_utils.map_structure_with_paths(convert, v)
      if converted and visited != converted:
        raise KeyError(
            "Only saw partial information from the dictionary for nested "
            "key '{}' in params_dict.  Entries provided: {}.  "
            "Entries required: {}"
            .format(k, sorted(converted), sorted(visited)))
    else:
      new_params[k] = convert(None, v)

  unvisited_params_keys = set(params_dict) - processed_params
  if unvisited_params_keys:
    raise ValueError(
        "params_dict had keys that were not part of value.params.  "
        "params_dict keys: {}, value.params processed keys: {}".format(
            sorted(params_dict.keys()), sorted(processed_params)))

  return Params(type_=value.type_, params=new_params)


def _check_no_tensors(parameters: Params):
  flat_params = tf.nest.flatten(parameters.params)
  for p in flat_params:
    if isinstance(p, Params):
      _check_no_tensors(p)
    if tf.is_tensor(p):
      raise TypeError(
          "Saw a `Tensor` value in parameters:\n  {}".format(parameters))


class DistributionSpecV2(object):
  """Describes a tfp.distribution.Distribution using nested parameters."""

  def __init__(self,
               event_shape: tf.TensorShape,
               dtype: tf.DType,
               parameters: Params):
    """Construct a `DistributionSpecV2` from a Distribution's properties.

    Note that the `parameters` used to create the spec should contain
    `tf.TypeSpec` objects instead of tensors.  We check for this.

    Args:
      event_shape: The distribution's `event_shape`.  This is the shape that
        `distribution.sample()` returns.  `distribution.sample(sample_shape)`
        returns tensors of shape `sample_shape + event_shape`.
      dtype: The distribution's `dtype`.
      parameters: The recursive parameters of the distribution, with
        tensors having directly been converted to `tf.TypeSpec` objects.

    Raises:
      TypeError: If for any entry `x` in `parameters`: `tf.is_tensor(x)`.
    """
    _check_no_tensors(parameters)
    self._event_shape = event_shape
    self._dtype = dtype
    self._parameters = parameters
    self._event_spec = tf.TensorSpec(shape=event_shape, dtype=dtype)

  @property
  def event_shape(self) -> tf.TensorShape:
    return self._event_shape

  @property
  def dtype(self) -> tf.DType:
    return self._dtype

  @property
  def event_spec(self) -> tf.TensorSpec:
    return self._event_spec

  @property
  def parameters(self) -> Params:
    return self._parameters

  def __eq__(self, other):
    return (isinstance(self, type(other))
            and self._event_shape == other._event_shape
            and self._dtype == other._dtype
            and self._parameters == other._parameters)

  def __str__(self):
    return ("<DistributionSpecV2: event_shape={}, dtype={}, parameters={}>"
            .format(self.event_shape, self.dtype, self.parameters))

  def __repr__(self):
    return str(self)


def assert_specs_are_compatible(
    network_output_spec: types.NestedTensorSpec,
    spec: types.NestedTensorSpec,
    message_prefix: str):
  """Checks that the output of `network.create_variables` matches a spec.

  Args:
    network_output_spec: The output of `network.create_variables`.
    spec: The spec we are matching to.
    message_prefix: The message prefix for error messages, used when the specs
      don't match.

  Raises:
    ValueError: If the specs don't match.
  """
  def to_event(s):
    return (s.event_spec if isinstance(s, DistributionSpecV2)
            else tensor_spec.from_spec(s))

  event_spec = tf.nest.map_structure(to_event, network_output_spec)

  nest_utils.assert_same_structure(
      event_spec,
      spec,
      message=("{}:\n{}\nvs.\n{}".format(message_prefix, event_spec, spec)))

  def compare_output_to_spec(s1, s2):
    if not s1.is_compatible_with(s2):
      raise ValueError("{}:\n{}\nvs.\n{}".format(message_prefix, event_spec,
                                                 spec))

  tf.nest.map_structure(compare_output_to_spec, event_spec, spec)
