# Copyright 2015 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.
# ==============================================================================
"""Wrapper allowing a stack of RNN cells to behave as a single cell."""


import functools

import tensorflow.compat.v2 as tf

from keras.src import backend
from keras.src.engine import base_layer
from keras.src.layers.rnn import rnn_utils
from keras.src.saving import serialization_lib
from keras.src.utils import generic_utils
from keras.src.utils import tf_utils

# isort: off
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util.tf_export import keras_export


@keras_export("keras.layers.StackedRNNCells")
class StackedRNNCells(base_layer.Layer):
    """Wrapper allowing a stack of RNN cells to behave as a single cell.

    Used to implement efficient stacked RNNs.

    Args:
      cells: List of RNN cell instances.

    Examples:

    ```python
    batch_size = 3
    sentence_max_length = 5
    n_features = 2
    new_shape = (batch_size, sentence_max_length, n_features)
    x = tf.constant(np.reshape(np.arange(30), new_shape), dtype = tf.float32)

    rnn_cells = [tf.keras.layers.LSTMCell(128) for _ in range(2)]
    stacked_lstm = tf.keras.layers.StackedRNNCells(rnn_cells)
    lstm_layer = tf.keras.layers.RNN(stacked_lstm)

    result = lstm_layer(x)
    ```
    """

    def __init__(self, cells, **kwargs):
        for cell in cells:
            if "call" not in dir(cell):
                raise ValueError(
                    "All cells must have a `call` method. "
                    f"Received cell without a `call` method: {cell}"
                )
            if "state_size" not in dir(cell):
                raise ValueError(
                    "All cells must have a `state_size` attribute. "
                    f"Received cell without a `state_size`: {cell}"
                )
        self.cells = cells
        # reverse_state_order determines whether the state size will be in a
        # reverse order of the cells' state. User might want to set this to True
        # to keep the existing behavior. This is only useful when use
        # RNN(return_state=True) since the state will be returned as the same
        # order of state_size.
        self.reverse_state_order = kwargs.pop("reverse_state_order", False)
        if self.reverse_state_order:
            logging.warning(
                "reverse_state_order=True in StackedRNNCells will soon "
                "be deprecated. Please update the code to work with the "
                "natural order of states if you rely on the RNN states, "
                "eg RNN(return_state=True)."
            )
        super().__init__(**kwargs)

    @property
    def state_size(self):
        return tuple(
            c.state_size
            for c in (
                self.cells[::-1] if self.reverse_state_order else self.cells
            )
        )

    @property
    def output_size(self):
        if getattr(self.cells[-1], "output_size", None) is not None:
            return self.cells[-1].output_size
        elif rnn_utils.is_multiple_state(self.cells[-1].state_size):
            return self.cells[-1].state_size[0]
        else:
            return self.cells[-1].state_size

    def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
        initial_states = []
        for cell in (
            self.cells[::-1] if self.reverse_state_order else self.cells
        ):
            get_initial_state_fn = getattr(cell, "get_initial_state", None)
            if get_initial_state_fn:
                initial_states.append(
                    get_initial_state_fn(
                        inputs=inputs, batch_size=batch_size, dtype=dtype
                    )
                )
            else:
                initial_states.append(
                    rnn_utils.generate_zero_filled_state_for_cell(
                        cell, inputs, batch_size, dtype
                    )
                )

        return tuple(initial_states)

    def call(self, inputs, states, constants=None, training=None, **kwargs):
        # Recover per-cell states.
        state_size = (
            self.state_size[::-1]
            if self.reverse_state_order
            else self.state_size
        )
        nested_states = tf.nest.pack_sequence_as(
            state_size, tf.nest.flatten(states)
        )

        # Call the cells in order and store the returned states.
        new_nested_states = []
        for cell, states in zip(self.cells, nested_states):
            states = states if tf.nest.is_nested(states) else [states]
            # TF cell does not wrap the state into list when there is only one
            # state.
            is_tf_rnn_cell = getattr(cell, "_is_tf_rnn_cell", None) is not None
            states = (
                states[0] if len(states) == 1 and is_tf_rnn_cell else states
            )
            if generic_utils.has_arg(cell.call, "training"):
                kwargs["training"] = training
            else:
                kwargs.pop("training", None)
            # Use the __call__ function for callable objects, eg layers, so that
            # it will have the proper name scopes for the ops, etc.
            cell_call_fn = cell.__call__ if callable(cell) else cell.call
            if generic_utils.has_arg(cell.call, "constants"):
                inputs, states = cell_call_fn(
                    inputs, states, constants=constants, **kwargs
                )
            else:
                inputs, states = cell_call_fn(inputs, states, **kwargs)
            new_nested_states.append(states)

        return inputs, tf.nest.pack_sequence_as(
            state_size, tf.nest.flatten(new_nested_states)
        )

    @tf_utils.shape_type_conversion
    def build(self, input_shape):
        if isinstance(input_shape, list):
            input_shape = input_shape[0]

        def get_batch_input_shape(batch_size, dim):
            shape = tf.TensorShape(dim).as_list()
            return tuple([batch_size] + shape)

        for cell in self.cells:
            if isinstance(cell, base_layer.Layer) and not cell.built:
                with backend.name_scope(cell.name):
                    cell.build(input_shape)
                    cell.built = True
            if getattr(cell, "output_size", None) is not None:
                output_dim = cell.output_size
            elif rnn_utils.is_multiple_state(cell.state_size):
                output_dim = cell.state_size[0]
            else:
                output_dim = cell.state_size
            batch_size = tf.nest.flatten(input_shape)[0]
            if tf.nest.is_nested(output_dim):
                input_shape = tf.nest.map_structure(
                    functools.partial(get_batch_input_shape, batch_size),
                    output_dim,
                )
                input_shape = tuple(input_shape)
            else:
                input_shape = tuple(
                    [batch_size] + tf.TensorShape(output_dim).as_list()
                )
        self.built = True

    def get_config(self):
        cells = []
        for cell in self.cells:
            cells.append(serialization_lib.serialize_keras_object(cell))
        config = {"cells": cells}
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))

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

        cells = []
        for cell_config in config.pop("cells"):
            cells.append(
                deserialize_layer(cell_config, custom_objects=custom_objects)
            )
        return cls(cells, **config)

