# Path: chatbot/components/base/_rnn.py
# class RNN(object):
#     """Base class for encoders/decoders. Has simple instance attributes and
#     an RNNCell object and getter.
#     """
# 
#     def __init__(self,
#                  state_size,
#                  embed_size,
#                  dropout_prob,
#                  num_layers,
#                  base_cell="GRUCell",
#                  state_wrapper=None):
#         """
#         Args:
#             state_size: number of units in underlying rnn cell.
#             embed_size: dimension size of word-embedding space.
#             dropout_prob: probability of a node being dropped.
#             num_layers: how many cells to include in the MultiRNNCell.
#             base_cell: (str) name of underling cell to use (e.g. 'GRUCell')
#             state_wrapper: allow states to store their wrapper class. See the
#                 wrapper method docstring below for more info.
#         """
#         self.state_size = state_size
#         self.embed_size = embed_size
#         self.num_layers = num_layers
#         self.dropout_prob = dropout_prob
#         self.base_cell = base_cell
#         self._wrapper = state_wrapper
# 
#     def get_cell(self, name):
#         """Returns a cell instance, defined by its name scope."""
#         with tf.name_scope(name, "get_cell"):
#             return Cell(state_size=self.state_size,
#                         num_layers=self.num_layers,
#                         dropout_prob=self.dropout_prob,
#                         base_cell=self.base_cell)
# 
#     def wrapper(self, state):
#         """Some RNN states are wrapped in namedtuples. 
#         (TensorFlow decision, definitely not mine...). 
#         
#         This is here for derived classes to specify their wrapper state. 
#         Some examples: LSTMStateTuple and AttentionWrapperState.
#         
#         Args:
#             state: tensor state tuple, will be unpacked into the wrapper tuple.
#         """
#         if self._wrapper is None:
#             return state
#         else:
#             return self._wrapper(*state)
# 
#     def __call__(self, *args):
#         raise NotImplemented


import tensorflow as tf
from tensorflow.contrib.rnn import GRUCell
from tensorflow.contrib.rnn import LSTMStateTuple, LSTMCell
from chatbot.components.base._rnn import RNN
from tensorflow.python.layers import core as layers_core



class BasicEncoder(RNN):
    """Encoder architecture that is defined by its cell running 
    inside dynamic_rnn.
    """

    def __call__(self, inputs, initial_state=None):
        """Run the inputs on the encoder and return the output(s).

        Args:
            inputs: Tensor with shape [batch_size, max_time, embed_size].
            initial_state: (optional) Tensor with shape [batch_size, state_size] 
                to initialize decoder cell.

        Returns:
            outputs: (only if return_sequence is True)
                     Tensor of shape [batch_size, max_time, state_size].
            state:   The final encoder state; shape [batch_size, state_size].
        """

        cell = self.get_cell("basic_enc_cell")
        _, state = tf.nn.dynamic_rnn(cell,
                                     inputs,
                                     initial_state=initial_state,
                                     dtype=tf.float32)
        return _, state



# Path: chatbot/components/base/_rnn.py
# class RNN(object):
#     """Base class for encoders/decoders. Has simple instance attributes and
#     an RNNCell object and getter.
#     """
# 
#     def __init__(self,
#                  state_size,
#                  embed_size,
#                  dropout_prob,
#                  num_layers,
#                  base_cell="GRUCell",
#                  state_wrapper=None):
#         """
#         Args:
#             state_size: number of units in underlying rnn cell.
#             embed_size: dimension size of word-embedding space.
#             dropout_prob: probability of a node being dropped.
#             num_layers: how many cells to include in the MultiRNNCell.
#             base_cell: (str) name of underling cell to use (e.g. 'GRUCell')
#             state_wrapper: allow states to store their wrapper class. See the
#                 wrapper method docstring below for more info.
#         """
#         self.state_size = state_size
#         self.embed_size = embed_size
#         self.num_layers = num_layers
#         self.dropout_prob = dropout_prob
#         self.base_cell = base_cell
#         self._wrapper = state_wrapper
# 
#     def get_cell(self, name):
#         """Returns a cell instance, defined by its name scope."""
#         with tf.name_scope(name, "get_cell"):
#             return Cell(state_size=self.state_size,
#                         num_layers=self.num_layers,
#                         dropout_prob=self.dropout_prob,
#                         base_cell=self.base_cell)
# 
#     def wrapper(self, state):
#         """Some RNN states are wrapped in namedtuples. 
#         (TensorFlow decision, definitely not mine...). 
#         
#         This is here for derived classes to specify their wrapper state. 
#         Some examples: LSTMStateTuple and AttentionWrapperState.
#         
#         Args:
#             state: tensor state tuple, will be unpacked into the wrapper tuple.
#         """
#         if self._wrapper is None:
#             return state
#         else:
#             return self._wrapper(*state)
# 
#     def __call__(self, *args):
#         raise NotImplemented


# Path: chatbot/components/encoders.py
import tensorflow as tf
from tensorflow.contrib.rnn import GRUCell
from tensorflow.contrib.rnn import LSTMStateTuple, LSTMCell
from chatbot.components.base._rnn import RNN
from tensorflow.python.layers import core as layers_core




class BasicEncoder(RNN):
    """Encoder architecture that is defined by its cell running 
    inside dynamic_rnn.
    """

    def __call__(self, inputs, initial_state=None):
        """Run the inputs on the encoder and return the output(s).

        Args:
            inputs: Tensor with shape [batch_size, max_time, embed_size].
            initial_state: (optional) Tensor with shape [batch_size, state_size] 
                to initialize decoder cell.

        Returns:
            outputs: (only if return_sequence is True)
                     Tensor of shape [batch_size, max_time, state_size].
            state:   The final encoder state; shape [batch_size, state_size].
        """

        cell = self.get_cell("basic_enc_cell")
        _, state = tf.nn.dynamic_rnn(cell,
                                     inputs,
                                     initial_state=initial_state,
                                     dtype=tf.float32)
        return _, state



class BidirectionalEncoder(RNN):