# Copyright 2019 The TensorFlow 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.
# ==============================================================================
"""Module implementing RNN wrappers."""


# Note that all the APIs under this module are exported as tf.nn.*. This is due
# to the fact that those APIs were from tf.nn.rnn_cell_impl. They are ported
# here to avoid the cyclic dependency issue for serialization. These APIs will
# probably be deprecated and removed in future since similar API is available in
# existing Keras RNN API.

import hashlib
import numbers
import sys
import types as python_types
import warnings

import tensorflow.compat.v2 as tf

from keras.layers.rnn import lstm
from keras.layers.rnn.abstract_rnn_cell import AbstractRNNCell
from keras.saving import serialization_lib
from keras.saving.legacy import serialization as legacy_serialization
from keras.utils import generic_utils
from keras.utils import tf_inspect

# isort: off
from tensorflow.python.util.tf_export import tf_export
from tensorflow.python.util.deprecation import deprecated


class _RNNCellWrapper(AbstractRNNCell):
    """Base class for cells wrappers V2 compatibility.

    This class along with `rnn_cell_impl._RNNCellWrapperV1` allows to define
    wrappers that are compatible with V1 and V2, and defines helper methods for
    this purpose.
    """

    def __init__(self, cell, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.cell = cell
        cell_call_spec = tf_inspect.getfullargspec(cell.call)
        self._call_spec.expects_training_arg = (
            "training" in cell_call_spec.args
        ) or (cell_call_spec.varkw is not None)

    def _call_wrapped_cell(self, inputs, state, cell_call_fn, **kwargs):
        """Calls the wrapped cell and performs the wrapping logic.

        This method is called from the wrapper's `call` or `__call__` methods.

        Args:
          inputs: A tensor with wrapped cell's input.
          state: A tensor or tuple of tensors with wrapped cell's state.
          cell_call_fn: Wrapped cell's method to use for step computation
            (cell's `__call__` or 'call' method).
          **kwargs: Additional arguments.

        Returns:
          A pair containing:
          - Output: A tensor with cell's output.
          - New state: A tensor or tuple of tensors with new wrapped cell's
            state.
        """
        raise NotImplementedError

    def call(self, inputs, state, **kwargs):
        """Runs the RNN cell step computation.

        When `call` is being used, we assume that the wrapper object has been
        built, and therefore the wrapped cells has been built via its `build`
        method and its `call` method can be used directly.

        This allows to use the wrapped cell and the non-wrapped cell
        equivalently when using `call` and `build`.

        Args:
          inputs: A tensor with wrapped cell's input.
          state: A tensor or tuple of tensors with wrapped cell's state.
          **kwargs: Additional arguments passed to the wrapped cell's `call`.

        Returns:
          A pair containing:

          - Output: A tensor with cell's output.
          - New state: A tensor or tuple of tensors with new wrapped cell's
            state.
        """
        return self._call_wrapped_cell(
            inputs, state, cell_call_fn=self.cell.call, **kwargs
        )

    def build(self, inputs_shape):
        """Builds the wrapped cell."""
        self.cell.build(inputs_shape)
        self.built = True

    @property
    def wrapped_cell(self):
        return self.cell

    @property
    def state_size(self):
        return self.cell.state_size

    @property
    def output_size(self):
        return self.cell.output_size

    def zero_state(self, batch_size, dtype):
        with tf.name_scope(type(self).__name__ + "ZeroState"):
            return self.cell.zero_state(batch_size, dtype)

    def get_config(self):
        config = {
            "cell": {
                "class_name": self.cell.__class__.__name__,
                "config": self.cell.get_config(),
            },
        }
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

    @classmethod
    def from_config(cls, config, custom_objects=None):
        config = config.copy()
        from keras.layers.serialization import deserialize as deserialize_layer

        cell = deserialize_layer(
            config.pop("cell"), custom_objects=custom_objects
        )
        return cls(cell, **config)


@deprecated(None, "Please use tf.keras.layers.RNN instead.")
@tf_export("nn.RNNCellDropoutWrapper", v1=[])
class DropoutWrapper(_RNNCellWrapper):
    """Operator adding dropout to inputs and outputs of the given cell."""

    def __init__(
        self,
        cell,
        input_keep_prob=1.0,
        output_keep_prob=1.0,
        state_keep_prob=1.0,
        variational_recurrent=False,
        input_size=None,
        dtype=None,
        seed=None,
        dropout_state_filter_visitor=None,
        **kwargs,
    ):
        """Create a cell with added input, state, and/or output dropout.

        If `variational_recurrent` is set to `True` (**NOT** the default
        behavior), then the same dropout mask is applied at every step, as
        described in: [A Theoretically Grounded Application of Dropout in
        Recurrent Neural Networks. Y. Gal, Z.
        Ghahramani](https://arxiv.org/abs/1512.05287).

        Otherwise a different dropout mask is applied at every time step.

        Note, by default (unless a custom `dropout_state_filter` is provided),
        the memory state (`c` component of any `LSTMStateTuple`) passing through
        a `DropoutWrapper` is never modified.  This behavior is described in the
        above article.

        Args:
          cell: an RNNCell, a projection to output_size is added to it.
          input_keep_prob: unit Tensor or float between 0 and 1, input keep
            probability; if it is constant and 1, no input dropout will be
            added.
          output_keep_prob: unit Tensor or float between 0 and 1, output keep
            probability; if it is constant and 1, no output dropout will be
            added.
          state_keep_prob: unit Tensor or float between 0 and 1, output keep
            probability; if it is constant and 1, no output dropout will be
            added.  State dropout is performed on the outgoing states of the
            cell. **Note** the state components to which dropout is applied when
            `state_keep_prob` is in `(0, 1)` are also determined by the argument
            `dropout_state_filter_visitor` (e.g. by default dropout is never
            applied to the `c` component of an `LSTMStateTuple`).
          variational_recurrent: Python bool.  If `True`, then the same dropout
            pattern is applied across all time steps per run call. If this
            parameter is set, `input_size` **must** be provided.
          input_size: (optional) (possibly nested tuple of) `TensorShape`
            objects containing the depth(s) of the input tensors expected to be
            passed in to the `DropoutWrapper`.  Required and used **iff**
            `variational_recurrent = True` and `input_keep_prob < 1`.
          dtype: (optional) The `dtype` of the input, state, and output tensors.
            Required and used **iff** `variational_recurrent = True`.
          seed: (optional) integer, the randomness seed.
          dropout_state_filter_visitor: (optional), default: (see below).
            Function that takes any hierarchical level of the state and returns
            a scalar or depth=1 structure of Python booleans describing which
            terms in the state should be dropped out.  In addition, if the
            function returns `True`, dropout is applied across this sublevel.
            If the function returns `False`, dropout is not applied across this
            entire sublevel.  Default behavior: perform dropout on all terms
            except the memory (`c`) state of `LSTMCellState` objects, and don't
            try to apply dropout to
            `TensorArray` objects:
            ```
            def dropout_state_filter_visitor(s):
              # Never perform dropout on the c state.
              if isinstance(s, LSTMCellState):
                return LSTMCellState(c=False, h=True)
              elif isinstance(s, TensorArray):
                return False
              return True
            ```
          **kwargs: dict of keyword arguments for base layer.

        Raises:
          TypeError: if `cell` is not an `RNNCell`, or `keep_state_fn` is
            provided but not `callable`.
          ValueError: if any of the keep_probs are not between 0 and 1.
        """
        if isinstance(cell, lstm.LSTMCell):
            raise ValueError(
                "keras LSTM cell does not work with DropoutWrapper. "
                "Please use LSTMCell(dropout=x, recurrent_dropout=y) "
                "instead."
            )
        super().__init__(cell, dtype=dtype, **kwargs)

        if dropout_state_filter_visitor is not None and not callable(
            dropout_state_filter_visitor
        ):
            raise TypeError(
                "dropout_state_filter_visitor must be callable. "
                f"Received: {dropout_state_filter_visitor}"
            )
        self._dropout_state_filter = (
            dropout_state_filter_visitor
            or _default_dropout_state_filter_visitor
        )
        with tf.name_scope("DropoutWrapperInit"):

            def tensor_and_const_value(v):
                tensor_value = tf.convert_to_tensor(v)
                const_value = tf.get_static_value(tensor_value)
                return (tensor_value, const_value)

            for prob, attr in [
                (input_keep_prob, "input_keep_prob"),
                (state_keep_prob, "state_keep_prob"),
                (output_keep_prob, "output_keep_prob"),
            ]:
                tensor_prob, const_prob = tensor_and_const_value(prob)
                if const_prob is not None:
                    if const_prob < 0 or const_prob > 1:
                        raise ValueError(
                            f"Parameter {attr} must be between 0 and 1. "
                            f"Received {const_prob}"
                        )
                    setattr(self, f"_{attr}", float(const_prob))
                else:
                    setattr(self, f"_{attr}", tensor_prob)

        # Set variational_recurrent, seed before running the code below
        self._variational_recurrent = variational_recurrent
        self._input_size = input_size
        self._seed = seed

        self._recurrent_input_noise = None
        self._recurrent_state_noise = None
        self._recurrent_output_noise = None

        if variational_recurrent:
            if dtype is None:
                raise ValueError(
                    "When variational_recurrent=True, dtype must be provided"
                )

            def convert_to_batch_shape(s):
                # Prepend a 1 for the batch dimension; for recurrent
                # variational dropout we use the same dropout mask for all
                # batch elements.
                return tf.concat(([1], tf.TensorShape(s).as_list()), 0)

            def batch_noise(s, inner_seed):
                shape = convert_to_batch_shape(s)
                return tf.random.uniform(shape, seed=inner_seed, dtype=dtype)

            if (
                not isinstance(self._input_keep_prob, numbers.Real)
                or self._input_keep_prob < 1.0
            ):
                if input_size is None:
                    raise ValueError(
                        "When variational_recurrent=True and input_keep_prob < "
                        "1.0 or is unknown, input_size must be provided"
                    )
                self._recurrent_input_noise = _enumerated_map_structure_up_to(
                    input_size,
                    lambda i, s: batch_noise(
                        s, inner_seed=self._gen_seed("input", i)
                    ),
                    input_size,
                )
            self._recurrent_state_noise = _enumerated_map_structure_up_to(
                cell.state_size,
                lambda i, s: batch_noise(
                    s, inner_seed=self._gen_seed("state", i)
                ),
                cell.state_size,
            )
            self._recurrent_output_noise = _enumerated_map_structure_up_to(
                cell.output_size,
                lambda i, s: batch_noise(
                    s, inner_seed=self._gen_seed("output", i)
                ),
                cell.output_size,
            )

    def _gen_seed(self, salt_prefix, index):
        if self._seed is None:
            return None
        salt = "%s_%d" % (salt_prefix, index)
        string = (str(self._seed) + salt).encode("utf-8")
        return int(hashlib.md5(string).hexdigest()[:8], 16) & 0x7FFFFFFF

    def _variational_recurrent_dropout_value(
        self, unused_index, value, noise, keep_prob
    ):
        """Performs dropout given the pre-calculated noise tensor."""
        # uniform [keep_prob, 1.0 + keep_prob)
        random_tensor = keep_prob + noise

        # 0. if [keep_prob, 1.0) and 1. if [1.0, 1.0 + keep_prob)
        binary_tensor = tf.floor(random_tensor)
        ret = tf.divide(value, keep_prob) * binary_tensor
        ret.set_shape(value.get_shape())
        return ret

    def _dropout(
        self,
        values,
        salt_prefix,
        recurrent_noise,
        keep_prob,
        shallow_filtered_substructure=None,
    ):
        """Decides whether to perform standard dropout or recurrent dropout."""

        if shallow_filtered_substructure is None:
            # Put something so we traverse the entire structure; inside the
            # dropout function we check to see if leafs of this are bool or not.
            shallow_filtered_substructure = values

        if not self._variational_recurrent:

            def dropout(i, do_dropout, v):
                if not isinstance(do_dropout, bool) or do_dropout:
                    return tf.nn.dropout(
                        v,
                        rate=1.0 - keep_prob,
                        seed=self._gen_seed(salt_prefix, i),
                    )
                else:
                    return v

            return _enumerated_map_structure_up_to(
                shallow_filtered_substructure,
                dropout,
                *[shallow_filtered_substructure, values],
            )
        else:

            def dropout(i, do_dropout, v, n):
                if not isinstance(do_dropout, bool) or do_dropout:
                    return self._variational_recurrent_dropout_value(
                        i, v, n, keep_prob
                    )
                else:
                    return v

            return _enumerated_map_structure_up_to(
                shallow_filtered_substructure,
                dropout,
                *[shallow_filtered_substructure, values, recurrent_noise],
            )

    def _call_wrapped_cell(self, inputs, state, cell_call_fn, **kwargs):
        """Runs the wrapped cell and applies dropout.

        Args:
          inputs: A tensor with wrapped cell's input.
          state: A tensor or tuple of tensors with wrapped cell's state.
          cell_call_fn: Wrapped cell's method to use for step computation
            (cell's `__call__` or 'call' method).
          **kwargs: Additional arguments.

        Returns:
          A pair containing:

          - Output: A tensor with cell's output.
          - New state: A tensor or tuple of tensors with new wrapped cell's
            state.
        """

        def _should_dropout(p):
            return (not isinstance(p, float)) or p < 1

        if _should_dropout(self._input_keep_prob):
            inputs = self._dropout(
                inputs,
                "input",
                self._recurrent_input_noise,
                self._input_keep_prob,
            )
        output, new_state = cell_call_fn(inputs, state, **kwargs)
        if _should_dropout(self._state_keep_prob):
            # Identify which subsets of the state to perform dropout on and
            # which ones to keep.
            shallow_filtered_substructure = (
                tf.__internal__.nest.get_traverse_shallow_structure(
                    self._dropout_state_filter, new_state
                )
            )
            new_state = self._dropout(
                new_state,
                "state",
                self._recurrent_state_noise,
                self._state_keep_prob,
                shallow_filtered_substructure,
            )
        if _should_dropout(self._output_keep_prob):
            output = self._dropout(
                output,
                "output",
                self._recurrent_output_noise,
                self._output_keep_prob,
            )
        return output, new_state

    def get_config(self):
        """Returns the config of the dropout wrapper."""
        config = {
            "input_keep_prob": self._input_keep_prob,
            "output_keep_prob": self._output_keep_prob,
            "state_keep_prob": self._state_keep_prob,
            "variational_recurrent": self._variational_recurrent,
            "input_size": self._input_size,
            "seed": self._seed,
        }
        if self._dropout_state_filter != _default_dropout_state_filter_visitor:
            (
                function,
                function_type,
                function_module,
            ) = _serialize_function_to_config(self._dropout_state_filter)
            config.update(
                {
                    "dropout_fn": function,
                    "dropout_fn_type": function_type,
                    "dropout_fn_module": function_module,
                }
            )
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

    @classmethod
    def from_config(cls, config, custom_objects=None):
        if "dropout_fn" in config:
            config = config.copy()
            dropout_state_filter = _parse_config_to_function(
                config,
                custom_objects,
                "dropout_fn",
                "dropout_fn_type",
                "dropout_fn_module",
            )
            config.pop("dropout_fn")
            config["dropout_state_filter_visitor"] = dropout_state_filter
        return super(DropoutWrapper, cls).from_config(
            config, custom_objects=custom_objects
        )


@deprecated(None, "Please use tf.keras.layers.RNN instead.")
@tf_export("nn.RNNCellResidualWrapper", v1=[])
class ResidualWrapper(_RNNCellWrapper):
    """RNNCell wrapper that ensures cell inputs are added to the outputs."""

    def __init__(self, cell, residual_fn=None, **kwargs):
        """Constructs a `ResidualWrapper` for `cell`.

        Args:
          cell: An instance of `RNNCell`.
          residual_fn: (Optional) The function to map raw cell inputs and raw
            cell outputs to the actual cell outputs of the residual network.
            Defaults to calling nest.map_structure on (lambda i, o: i + o),
            inputs and outputs.
          **kwargs: dict of keyword arguments for base layer.
        """
        super().__init__(cell, **kwargs)
        self._residual_fn = residual_fn

    def _call_wrapped_cell(self, inputs, state, cell_call_fn, **kwargs):
        """Run the cell and apply the residual_fn.

        Args:
          inputs: cell inputs.
          state: cell state.
          cell_call_fn: Wrapped cell's method to use for step computation
            (cell's `__call__` or 'call' method).
          **kwargs: Additional arguments passed to the wrapped cell's `call`.

        Returns:
          Tuple of cell outputs and new state.

        Raises:
          TypeError: If cell inputs and outputs have different structure (type).
          ValueError: If cell inputs and outputs have different structure
            (value).
        """
        outputs, new_state = cell_call_fn(inputs, state, **kwargs)

        # Ensure shapes match
        def assert_shape_match(inp, out):
            inp.get_shape().assert_is_compatible_with(out.get_shape())

        def default_residual_fn(inputs, outputs):
            tf.nest.assert_same_structure(inputs, outputs)
            tf.nest.map_structure(assert_shape_match, inputs, outputs)
            return tf.nest.map_structure(
                lambda inp, out: inp + out, inputs, outputs
            )

        res_outputs = (self._residual_fn or default_residual_fn)(
            inputs, outputs
        )
        return (res_outputs, new_state)

    def get_config(self):
        """Returns the config of the residual wrapper."""
        if self._residual_fn is not None:
            (
                function,
                function_type,
                function_module,
            ) = _serialize_function_to_config(self._residual_fn)
            config = {
                "residual_fn": function,
                "residual_fn_type": function_type,
                "residual_fn_module": function_module,
            }
        else:
            config = {}
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

    @classmethod
    def from_config(cls, config, custom_objects=None):
        if "residual_fn" in config:
            config = config.copy()
            residual_function = _parse_config_to_function(
                config,
                custom_objects,
                "residual_fn",
                "residual_fn_type",
                "residual_fn_module",
            )
            config["residual_fn"] = residual_function
        return super(ResidualWrapper, cls).from_config(
            config, custom_objects=custom_objects
        )


@deprecated(None, "Please use tf.keras.layers.RNN instead.")
@tf_export("nn.RNNCellDeviceWrapper", v1=[])
class DeviceWrapper(_RNNCellWrapper):
    """Operator that ensures an RNNCell runs on a particular device."""

    def __init__(self, cell, device, **kwargs):
        """Construct a `DeviceWrapper` for `cell` with device `device`.

        Ensures the wrapped `cell` is called with `tf.device(device)`.

        Args:
          cell: An instance of `RNNCell`.
          device: A device string or function, for passing to `tf.device`.
          **kwargs: dict of keyword arguments for base layer.
        """
        super().__init__(cell, **kwargs)
        self._device = device

    def zero_state(self, batch_size, dtype):
        with tf.name_scope(type(self).__name__ + "ZeroState"):
            with tf.compat.v1.device(self._device):
                return self.cell.zero_state(batch_size, dtype)

    def _call_wrapped_cell(self, inputs, state, cell_call_fn, **kwargs):
        """Run the cell on specified device."""
        with tf.compat.v1.device(self._device):
            return cell_call_fn(inputs, state, **kwargs)

    def get_config(self):
        config = {"device": self._device}
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))


def _serialize_function_to_config(function):
    """Serialize the function for get_config()."""
    if isinstance(function, python_types.LambdaType):
        output = generic_utils.func_dump(function)
        output_type = "lambda"
        module = function.__module__
    elif callable(function):
        output = function.__name__
        output_type = "function"
        module = function.__module__
    else:
        raise ValueError(
            f"Unrecognized function type for input: {type(function)}"
        )

    return output, output_type, module


def _parse_config_to_function(
    config,
    custom_objects,
    func_attr_name,
    func_type_attr_name,
    module_attr_name,
):
    """Reconstruct the function from the config."""
    globs = globals()
    module = config.pop(module_attr_name, None)
    if module in sys.modules:
        globs.update(sys.modules[module].__dict__)
    elif module is not None:
        # Note: we don't know the name of the function if it's a lambda.
        warnings.warn(
            "{} is not loaded, but a layer uses it. "
            "It may cause errors.".format(module),
            UserWarning,
            stacklevel=2,
        )
    if custom_objects:
        globs.update(custom_objects)
    function_type = config.pop(func_type_attr_name)
    if function_type == "function":
        # Simple lookup in custom objects
        function = legacy_serialization.deserialize_keras_object(
            config[func_attr_name],
            custom_objects=custom_objects,
            printable_module_name="function in wrapper",
        )
    elif function_type == "lambda":
        if serialization_lib.in_safe_mode():
            raise ValueError(
                "Requested the deserialization of a layer with a "
                "Python `lambda` inside it. "
                "This carries a potential risk of arbitrary code execution "
                "and thus it is disallowed by default. If you trust the "
                "source of the saved model, you can pass `safe_mode=False` to "
                "the loading function in order to allow "
                "`lambda` loading."
            )
        # Unsafe deserialization from bytecode
        function = generic_utils.func_load(config[func_attr_name], globs=globs)
    else:
        raise TypeError(
            f"Unknown function type received: {function_type}. "
            "Expected types are ['function', 'lambda']"
        )
    return function


def _default_dropout_state_filter_visitor(substate):
    return not isinstance(substate, tf.TensorArray)


def _enumerated_map_structure_up_to(shallow_structure, map_fn, *args, **kwargs):
    ix = [0]

    def enumerated_fn(*inner_args, **inner_kwargs):
        r = map_fn(ix[0], *inner_args, **inner_kwargs)
        ix[0] += 1
        return r

    return tf.__internal__.nest.map_structure_up_to(
        shallow_structure, enumerated_fn, *args, **kwargs
    )
