"""Define bridges: logic of passing the encoder state to the decoder."""

import abc

import tensorflow as tf


def assert_state_is_compatible(expected_state, state):
  """Asserts that states are compatible.

  Args:
    expected_state: The reference state.
    state: The state that must be compatible with :obj:`expected_state`.

  Raises:
    ValueError: if the states are incompatible.
  """
  # Check structure compatibility.
  tf.nest.assert_same_structure(expected_state, state)

  # Check shape compatibility.
  expected_state_flat = tf.nest.flatten(expected_state)
  state_flat = tf.nest.flatten(state)

  for x, y in zip(expected_state_flat, state_flat):
    if tf.is_tensor(x):
      expected_depth = x.shape[-1]
      depth = y.shape[-1]
      if depth != expected_depth:
        raise ValueError("Tensor in state has shape %s which is incompatible "
                         "with the target shape %s" % (y.shape, x.shape))


class Bridge(tf.keras.layers.Layer):
  """Base class for bridges."""

  def __call__(self, encoder_state, decoder_zero_state):  # pylint: disable=arguments-differ
    """Returns the initial decoder state.

    Args:
      encoder_state: The encoder state.
      decoder_zero_state: The default decoder state.

    Returns:
      The decoder initial state.
    """
    return super(Bridge, self).__call__([encoder_state, decoder_zero_state])

  @abc.abstractmethod
  def call(self, states):  # pylint: disable=arguments-differ
    raise NotImplementedError()


class CopyBridge(Bridge):
  """A bridge that passes the encoder state as is."""

  def call(self, states):
    encoder_state, decoder_state = states
    assert_state_is_compatible(encoder_state, decoder_state)
    flat_encoder_state = tf.nest.flatten(encoder_state)
    return tf.nest.pack_sequence_as(decoder_state, flat_encoder_state)


class ZeroBridge(Bridge):
  """A bridge that does not pass information from the encoder."""

  def call(self, states):
    # Simply return the default decoder state.
    return states[1]


class DenseBridge(Bridge):
  """A bridge that applies a parameterized linear transformation from the
  encoder state to the decoder state size.
  """

  def __init__(self, activation=None):
    """Initializes the bridge.

    Args:
      activation: Activation function (a callable).
        Set it to ``None`` to maintain a linear activation.
    """
    super(DenseBridge, self).__init__()
    self.activation = activation
    self.decoder_state_sizes = None
    self.linear = None

  def build(self, input_shape):
    decoder_shape = input_shape[1]
    self.decoder_state_sizes = [
        shape[-1] for shape in tf.nest.flatten(decoder_shape)]
    self.linear = tf.keras.layers.Dense(
        sum(self.decoder_state_sizes), activation=self.activation)

  def call(self, states):
    encoder_state, decoder_state = states
    encoder_state_flat = tf.nest.flatten(encoder_state)
    encoder_state_single = tf.concat(encoder_state_flat, 1)
    transformed = self.linear(encoder_state_single)
    splitted = tf.split(transformed, self.decoder_state_sizes, axis=1)
    return tf.nest.pack_sequence_as(decoder_state, splitted)
