"""Network units wrapping TensorFlows' tf.contrib.rnn cells.

Please put all wrapping logic for tf.contrib.rnn in this module; this will help
collect common subroutines that prove useful.
"""

import abc

import tensorflow as tf

from dragnn.python import network_units as dragnn
from syntaxnet.util import check


class BaseLSTMNetwork(dragnn.NetworkUnitInterface):
  """Base class for wrapped LSTM networks.

  This LSTM network unit supports multiple layers with layer normalization.
  Because it is imported from tf.contrib.rnn, we need to capture the created
  variables during initialization time.

  Layers:
    ...subclass-specific layers...
    last_layer: Alias for the activations of the last hidden layer.
    logits: Logits associated with component actions.
  """

  def __init__(self, component):
    """Initializes the LSTM base class.

    Parameters used:
      hidden_layer_sizes: Comma-delimited number of hidden units for each layer.
      input_dropout_rate (-1.0): Input dropout rate for each layer.  If < 0.0,
          use the global |dropout_rate| hyperparameter.
      recurrent_dropout_rate (0.8): Recurrent dropout rate.  If < 0.0, use the
          global |recurrent_dropout_rate| hyperparameter.
      layer_norm (True): Whether or not to use layer norm.

    Hyperparameters used:
      dropout_rate: Input dropout rate.
      recurrent_dropout_rate: Recurrent dropout rate.

    Args:
      component: parent ComponentBuilderBase object.
    """
    self._attrs = dragnn.get_attrs_with_defaults(
        component.spec.network_unit.parameters,
        defaults={
            'layer_norm': True,
            'input_dropout_rate': -1.0,
            'recurrent_dropout_rate': 0.8,
            'hidden_layer_sizes': '256',
        })

    self._hidden_layer_sizes = map(int,
                                   self._attrs['hidden_layer_sizes'].split(','))

    self._input_dropout_rate = self._attrs['input_dropout_rate']
    if self._input_dropout_rate < 0.0:
      self._input_dropout_rate = component.master.hyperparams.dropout_rate

    self._recurrent_dropout_rate = self._attrs['recurrent_dropout_rate']
    if self._recurrent_dropout_rate < 0.0:
      self._recurrent_dropout_rate = (
          component.master.hyperparams.recurrent_dropout_rate)
    if self._recurrent_dropout_rate < 0.0:
      self._recurrent_dropout_rate = component.master.hyperparams.dropout_rate

    tf.logging.info('[%s] input_dropout_rate=%s recurrent_dropout_rate=%s',
                    component.name, self._input_dropout_rate,
                    self._recurrent_dropout_rate)

    layers, context_layers = self.create_hidden_layers(component,
                                                       self._hidden_layer_sizes)
    last_layer_dim = layers[-1].dim
    layers.append(
        dragnn.Layer(component, name='last_layer', dim=last_layer_dim))
    layers.append(
        dragnn.Layer(component, name='logits', dim=component.num_actions))

    # Provide initial layers and context layers, so the base class constructor
    # can safely use accessors like get_layer_size().
    super(BaseLSTMNetwork, self).__init__(
        component, init_layers=layers, init_context_layers=context_layers)

    # Allocate parameters for the softmax.
    self._params.append(
        tf.get_variable(
            'weights_softmax', [last_layer_dim, component.num_actions],
            initializer=tf.random_normal_initializer(
                stddev=1e-4, seed=self._seed)))
    self._params.append(
        tf.get_variable(
            'bias_softmax', [component.num_actions],
            initializer=tf.zeros_initializer()))

  def get_logits(self, network_tensors):
    """Returns the logits for prediction."""
    return network_tensors[self.get_layer_index('logits')]

  @abc.abstractmethod
  def create_hidden_layers(self, component, hidden_layer_sizes):
    """Creates hidden network layers.

    Args:
      component: Parent ComponentBuilderBase object.
      hidden_layer_sizes: List of requested hidden layer activation sizes.

    Returns:
      layers: List of layers created by this network.
      context_layers: List of context layers created by this network.
    """
    pass

  def _append_base_layers(self, hidden_layers):
    """Appends layers defined by the base class to the |hidden_layers|."""
    last_layer = hidden_layers[-1]

    # TODO(googleuser): Uncomment the version that uses component.get_variable()
    # and delete the uses of tf.get_variable().
    # logits = tf.nn.xw_plus_b(last_layer,
    #                          self._component.get_variable('weights_softmax'),
    #                          self._component.get_variable('bias_softmax'))
    logits = tf.nn.xw_plus_b(last_layer,
                             tf.get_variable('weights_softmax'),
                             tf.get_variable('bias_softmax'))
    return hidden_layers + [last_layer, logits]

  def _create_cell(self, num_units, during_training):
    """Creates a single LSTM cell, possibly with dropout.

    Requires that BaseLSTMNetwork.__init__() was called.

    Args:
      num_units: Number of hidden units in the cell.
      during_training: Whether to create a cell for training (vs inference).

    Returns:
      A RNNCell of the requested size, possibly with dropout.
    """
    # No dropout in inference mode.
    if not during_training:
      return tf.contrib.rnn.LayerNormBasicLSTMCell(
          num_units, layer_norm=self._attrs['layer_norm'], reuse=True)

    # Otherwise, apply dropout to inputs and recurrences.
    cell = tf.contrib.rnn.LayerNormBasicLSTMCell(
        num_units,
        dropout_keep_prob=self._recurrent_dropout_rate,
        layer_norm=self._attrs['layer_norm'])
    cell = tf.contrib.rnn.DropoutWrapper(
        cell, input_keep_prob=self._input_dropout_rate)
    return cell

  def _create_train_cells(self):
    """Creates a list of LSTM cells for training."""
    return [
        self._create_cell(num_units, during_training=True)
        for num_units in self._hidden_layer_sizes
    ]

  def _create_inference_cells(self):
    """Creates a list of LSTM cells for inference."""
    return [
        self._create_cell(num_units, during_training=False)
        for num_units in self._hidden_layer_sizes
    ]

  def _capture_variables_as_params(self, function):
    """Captures variables created by a function in |self._params|.

    Args:
      function: Function whose variables should be captured.  The function
          should take one argument, its enclosing variable scope.
    """
    created_vars = {}

    def _custom_getter(getter, *args, **kwargs):
      """Calls the real getter and captures its result in |created_vars|."""
      real_variable = getter(*args, **kwargs)
      created_vars[real_variable.name] = real_variable
      return real_variable

    with tf.variable_scope(
        'cell', reuse=None, custom_getter=_custom_getter) as scope:
      function(scope)
    self._params.extend(created_vars.values())

  def _apply_with_captured_variables(self, function):
    """Applies a function using previously-captured variables.

    Args:
      function: Function to apply using captured variables.  The function
          should take one argument, its enclosing variable scope.

    Returns:
      Results of function application.
    """

    def _custom_getter(getter, *args, **kwargs):
      """Retrieves the normal or moving-average variables."""
      return self._component.get_variable(var_params=getter(*args, **kwargs))

    with tf.variable_scope(
        'cell', reuse=True, custom_getter=_custom_getter) as scope:
      return function(scope)


class LayerNormBasicLSTMNetwork(BaseLSTMNetwork):
  """Wrapper around tf.contrib.rnn.LayerNormBasicLSTMCell.

  Features:
    All inputs are concatenated.

  Subclass-specific layers:
    state_c_<n>: Cell states for the <n>'th LSTM layer (0-origin).
    state_h_<n>: Hidden states for the <n>'th LSTM layer (0-origin).
  """

  def __init__(self, component):
    """Sets up context and output layers, as well as a final softmax."""
    super(LayerNormBasicLSTMNetwork, self).__init__(component)

    # Wrap lists of training and inference sub-cells into multi-layer RNN cells.
    # Note that a |MultiRNNCell| state is a tuple of per-layer sub-states.
    self._train_cell = tf.contrib.rnn.MultiRNNCell(self._create_train_cells())
    self._inference_cell = tf.contrib.rnn.MultiRNNCell(
        self._create_inference_cells())

    def _cell_closure(scope):
      """Applies the LSTM cell to placeholder inputs and state."""
      placeholder_inputs = tf.placeholder(
          dtype=tf.float32, shape=(1, self._concatenated_input_dim))

      placeholder_substates = []
      for num_units in self._hidden_layer_sizes:
        placeholder_substate = tf.contrib.rnn.LSTMStateTuple(
            tf.placeholder(dtype=tf.float32, shape=(1, num_units)),
            tf.placeholder(dtype=tf.float32, shape=(1, num_units)))
        placeholder_substates.append(placeholder_substate)
      placeholder_state = tuple(placeholder_substates)

      self._train_cell(
          inputs=placeholder_inputs, state=placeholder_state, scope=scope)

    self._capture_variables_as_params(_cell_closure)

  def create_hidden_layers(self, component, hidden_layer_sizes):
    """See base class."""
    # Construct the layer meta info for the DRAGNN builder. Note that the order
    # of h and c are reversed compared to the vanilla DRAGNN LSTM cell, as
    # this is the standard in tf.contrib.rnn.
    #
    # NB: The h activations of the last LSTM must be the last layer, in order
    # for _append_base_layers() to work.
    layers = []
    for index, num_units in enumerate(hidden_layer_sizes):
      layers.append(
          dragnn.Layer(component, name='state_c_%d' % index, dim=num_units))
      layers.append(
          dragnn.Layer(component, name='state_h_%d' % index, dim=num_units))