# coding=utf-8
# Copyright 2020 The Tensor2Tensor Authors.
#
# 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.

"""Modalities, which specify a feature's domain.

T2TModel applies a default transformation to each feature according to its
modality. Override them by specifying a model's
hparams.{bottom,loss,top,weights_fn}.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from six.moves import range  # pylint: disable=redefined-builtin

from tensor2tensor.layers import common_attention
from tensor2tensor.layers import common_audio
from tensor2tensor.layers import common_image_attention as cia
from tensor2tensor.layers import common_layers
from tensor2tensor.layers import common_video
from tensor2tensor.layers import discretization

import tensorflow.compat.v1 as tf
import tensorflow_probability as tfp


class ModalityType(object):
  """Types of modalities."""

  AUDIO = "audio"
  AUDIO_SPECTRAL = "audio_spectral"
  CLASS_LABEL = "class_label"
  CTC_SYMBOL = "ctc_symbol"  # symbol with CTC loss
  GENERIC_L2_LOSS = "generic_l2"  # identity modality with L2 loss
  IDENTITY = "identity"  # identity top and bottom
  IDENTITY_SYMBOL = "identity_symbol"  # symbol with identity top and bottom
  IMAGE = "image"
  # images using channel compression for generation
  IMAGE_CHANNEL_BOTTOM_IDENTITY = "image_channel_bottom_identity"
  # images using channel compression for generation
  IMAGE_CHANNEL_COMPRESS = "image_channel_compress"
  IMAGE_CHANNEL_EMBEDDINGS_BOTTOM = "image_channel_embeddings_bottom"
  MULTI_LABEL = "multi_label"
  ONE_HOT_CLASS_LABEL = "one_hot_class_label"
  REAL = "real"  # real vectors
  REAL_L2_LOSS = "real_l2"  # real vectors with L2 as loss
  # real vectors with log Poisson regression loss
  REAL_LOG_POISSON_LOSS = "real_log_poisson"
  SIGMOID_CLASS_LABEL = "sigmoid_class_label"  # sigmoid cross-entropy loss
  # sigmoid cross-entropy applied on max-pooling over timesteps
  SIGMOID_MAX_POOLING_CLASS_LABEL = "sigmoid_max_pooling_class_label"
  # softmax cross-entropy applied on average-pooling over timesteps
  SOFTMAX_AVERAGE_POOLING_CLASS_LABEL = "softmax_average_pooling_class_label"
  # softmax cross-entropy applied on last-timestep encoding
  SOFTMAX_LAST_TIMESTEP_CLASS_LABEL = "softmax_last_timestep_class_label"
  # softmax cross-entropy applied on max-pooling over timesteps
  SOFTMAX_MAX_POOLING_CLASS_LABEL = "softmax_max_pooling_class_label"
  SPEECH_RECOGNITION = "speech_recognition"
  SYMBOL = "symbol"
  SYMBOL_WEIGHTS_ALL = "symbol_weights_all"  # symbol for features w/o 0-padding
  SYMBOL_ONE_HOT = "symbol_one_hot"  # symbol with one hot as embeddings
  VIDEO = "video"
  VIDEO_BITWISE = "video_bitwise"  # video where bottom embeds pixels bitwise
  VIDEO_IDENTITY = "video_identity"  # video with identity top and bottom
  VIDEO_L1 = "video_l1"  # video with L2 loss
  VIDEO_L2 = "video_l2"  # video with L1 loss
  # video with L1 loss and raw input (sequences of frames)
  VIDEO_L1_RAW = "video_l1_raw"
  # video with L2 loss and raw input (sequences of frames)
  VIDEO_L2_RAW = "video_l2_raw"
  # video with pixel noise on input during training
  VIDEO_PIXEL_NOISE = "video_pixel_noise"

  @staticmethod
  def get_choices():
    return [
        ModalityType.AUDIO,
        ModalityType.AUDIO_SPECTRAL,
        ModalityType.CLASS_LABEL,
        ModalityType.CTC_SYMBOL,
        ModalityType.GENERIC_L2_LOSS,
        ModalityType.IDENTITY,
        ModalityType.IDENTITY_SYMBOL,
        ModalityType.IMAGE,
        ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
        ModalityType.IMAGE_CHANNEL_COMPRESS,
        ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM,
        ModalityType.MULTI_LABEL,
        ModalityType.ONE_HOT_CLASS_LABEL,
        ModalityType.REAL,
        ModalityType.REAL_L2_LOSS,
        ModalityType.REAL_LOG_POISSON_LOSS,
        ModalityType.SIGMOID_CLASS_LABEL,
        ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
        ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
        ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
        ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL,
        ModalityType.SPEECH_RECOGNITION,
        ModalityType.SYMBOL,
        ModalityType.SYMBOL_ONE_HOT,
        ModalityType.SYMBOL_WEIGHTS_ALL,
        ModalityType.VIDEO,
        ModalityType.VIDEO_BITWISE,
        ModalityType.VIDEO_IDENTITY,
        ModalityType.VIDEO_L1,
        ModalityType.VIDEO_L2,
        ModalityType.VIDEO_L1_RAW,
        ModalityType.VIDEO_L2_RAW,
        ModalityType.VIDEO_PIXEL_NOISE,
    ]


# Bottom transformations, applied to all features


def audio_bottom(x, model_hparams, vocab_size):
  """Transform input from data space to model space.

  Args:
    x: A Tensor with shape [batch, ...]
    model_hparams: HParams, model hyperparmeters.
    vocab_size: int, vocabulary size.

  Returns:
    body_input: A Tensor with shape [batch, ?, ?,
      model_hparams.hidden_size].
  """
  del vocab_size  # unused arg
  inputs = x
  with tf.variable_scope("audio_modality"):
    # TODO(aidangomez): Will need to sort out a better audio pipeline
    def xnet_resblock(x, filters, res_relu, name):
      """Xception block."""
      with tf.variable_scope(name):
        # Typically audio samples are >100k samples in length and have a width
        # of 2 or 4. Mono audio has a single channel while stereo has 2.
        y = common_layers.separable_conv_block(
            x,
            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
            first_relu=True,
            padding="SAME",
            force2d=True,
            name="sep_conv_block")
        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 2))
        return y + common_layers.conv_block(
            x,
            filters, [((1, 1), (1, 1))],
            padding="SAME",
            strides=(2, 2),
            first_relu=res_relu,
            force2d=True,
            name="res_conv0")

    x = tf.to_float(inputs) / 255.
    x.set_shape([None, None, None, 1])
    for i in range(model_hparams.audio_compression):
      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
    return xnet_resblock(x,
                         model_hparams.hidden_size,
                         False,
                         "compress_block_final")


def audio_spectral_bottom(x, model_hparams, vocab_size):
  """Transform input from data space to model space.

  Args:
    x: A Tensor with shape [batch, ...]
    model_hparams: HParams, model hyperparmeters.
    vocab_size: int, vocabulary size.

  Returns:
    body_input: A Tensor with shape [batch, ?, ?,
      model_hparams.hidden_size].
  """
  del vocab_size  # unused arg
  inputs = x
  with tf.variable_scope("audio_spectral_modality"):
    # TODO(aidangomez): Will need to sort out a better audio pipeline
    def xnet_resblock(x, filters, res_relu, name):
      """Xception-like block."""
      with tf.variable_scope(name):
        # We only stride along the length dimension to preserve the spectral
        # bins (which are tiny in dimensionality relative to length)
        y = common_layers.separable_conv_block(
            x,
            filters, [((1, 1), (3, 3)), ((1, 1), (3, 3))],
            first_relu=True,
            padding="SAME",
            force2d=True,
            name="sep_conv_block")
        y = common_layers.pool(y, (3, 3), "MAX", "SAME", strides=(2, 1))
        return y + common_layers.conv_block(
            x,
            filters, [((1, 1), (1, 1))],
            padding="SAME",
            strides=(2, 1),
            first_relu=res_relu,
            force2d=True,
            name="res_conv0")

    # Bitcast back from int32
    x = tf.bitcast(inputs, tf.float32)
    x.set_shape([None, None, None, 1])
    for i in range(model_hparams.audio_compression):
      x = xnet_resblock(x, 2**(i + 1), True, "compress_block_%d" % i)
    return xnet_resblock(x,
                         model_hparams.hidden_size,
                         False,
                         "compress_block_final")


def class_label_bottom(x, model_hparams, vocab_size):
  with tf.variable_scope("class_label_modality_%d_%d" % (
      vocab_size, model_hparams.hidden_size)):
    multiplier = 1.0
    if model_hparams.multiply_embedding_mode == "sqrt_depth":
      multiplier = model_hparams.hidden_size**0.5
    return common_layers.embedding(x,
                                   vocab_size,
                                   model_hparams.hidden_size,
                                   multiplier=multiplier)


def class_label_targets_bottom(x, model_hparams, vocab_size):
  with tf.variable_scope("class_label_modality_%d_%d" % (
      vocab_size, model_hparams.hidden_size)):
    return tf.zeros([common_layers.shape_list(x)[0],
                     1,
                     1,
                     model_hparams.hidden_size])


def identity_bottom(x, model_hparams, vocab_size):
  del model_hparams, vocab_size  # unused arg
  return tf.to_float(x)


def image_bottom(x, model_hparams, vocab_size):
  del model_hparams, vocab_size  # unused arg
  with tf.variable_scope("image_modality"):
    if not tf.executing_eagerly():
      tf.summary.image(
          "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
    return tf.to_float(x)


def image_targets_bottom(x, model_hparams, vocab_size):
  """Bottom transformation for target images."""
  pixel_embedding_size = 64
  inputs = x
  with tf.variable_scope("image_modality"):
    if not tf.executing_eagerly():
      tf.summary.image(
          "targets_bottom",
          common_layers.tpu_safe_image_summary(inputs),
          max_outputs=1)
    inputs_shape = common_layers.shape_list(inputs)
    if len(inputs_shape) != 4:
      raise ValueError("Assuming images given as int tensors in the format "
                       "[batch, height, width, channels] (256 values).")
    # We embed each of 256=vocab_size possible pixel values.
    embedding_var = tf.get_variable(
        "pixel_embedding",
        [vocab_size, pixel_embedding_size])
    hot_inputs = tf.one_hot(tf.to_int32(inputs), vocab_size)
    hot_inputs = tf.reshape(hot_inputs, [-1, vocab_size])
    embedded = tf.matmul(hot_inputs, embedding_var)
    # Let's now merge all channels that were embedded into a single vector.
    merged_size = pixel_embedding_size * inputs_shape[3]
    embedded = tf.reshape(embedded, inputs_shape[:3] + [merged_size])
    merged = tf.layers.dense(
        embedded,
        model_hparams.hidden_size,
        name="merge_pixel_embedded_channels")
    return merged


def _image_channel_compress_bottom(inputs, model_hparams, name="bottom"):
  """Compresses channel-wise input pixels into whole pixel representions.

  Perform conversion of RGB pixel values to a real number in the range -1 to
  1. This combines pixel channels to form a representation of shape
  [img_len, img_len].

  Args:
    inputs: Tensor representing RGB pixel intensities as integers, of shape
      [batch, img_len, img_len, channels].
    model_hparams: HParams, model hyperparmeters.
    name: string, scope.

  Returns:
    body_input: Tensor of shape
      [batch, img_len, img_len, model_hparams.hidden_size].
  """
  num_channels = 3
  with tf.variable_scope(name):
    inputs = tf.to_float(inputs)
    hp = model_hparams
    if hp.mode != tf.estimator.ModeKeys.PREDICT:
      tf.summary.image(
          "inputs",
          common_layers.tpu_safe_image_summary(inputs),
          max_outputs=2)
    inputs = common_layers.convert_rgb_to_symmetric_real(inputs)

    # Reshape inputs to apply convolutions across [img_len, img_len*channels].
    inputs_shape = common_layers.shape_list(inputs)
    inputs = tf.reshape(
        inputs, [-1, inputs_shape[1], inputs_shape[2] * inputs_shape[3], 1])

    # Compress RGB intensities for each pixel using a convolution.
    outputs = tf.layers.conv2d(
        inputs,
        model_hparams.hidden_size,
        kernel_size=(1, num_channels),
        padding="VALID",
        strides=(1, num_channels),
        activation=tf.nn.relu,
        name="conv_input")
    return outputs


def image_channel_compress_bottom(x, model_hparams, vocab_size):
  del vocab_size  # unused arg
  return _image_channel_compress_bottom(x, model_hparams, "input_bottom")


def image_channel_compress_targets_bottom(x, model_hparams, vocab_size):
  del vocab_size  # unused arg
  return _image_channel_compress_bottom(x, model_hparams, "output_bottom")


def image_channel_embeddings_bottom(x, model_hparams, vocab_size):
  """Bottom transformation for image targets."""
  del vocab_size  # unused arg
  inputs = tf.to_int32(x)
  io_depth = model_hparams.num_channels
  tshape = common_layers.shape_list(inputs)
  hidden_size = model_hparams.hidden_size
  target_embeddings = cia.get_channel_embeddings(
      io_depth, inputs, hidden_size, "input_bottom")
  return tf.reshape(target_embeddings,
                    [tshape[0], tshape[1], tshape[2] * io_depth, hidden_size])


def make_targets_bottom(bottom):
  def targets_bottom(x, model_hparams, vocab_size):
    with tf.variable_scope("targets_bottom"):
      return bottom(x, model_hparams, vocab_size)
  return targets_bottom


def real_bottom(x, model_hparams, vocab_size):
  del vocab_size  # unused arg
  with tf.variable_scope("real"):
    return tf.layers.dense(
        tf.to_float(x), model_hparams.hidden_size, name="bottom")


def speech_recognition_bottom(x, model_hparams, vocab_size):
  """Use batchnorm instead of CMVN and shorten the stft with strided convs.

  Args:
    x: float32 tensor with shape [batch_size, len, 1, freqs * channels]
    model_hparams: HParams, model hyperparmeters.
    vocab_size: int, vocabulary size.

  Returns:
    float32 tensor with shape [batch_size, shorter_len, 1, hidden_size]
  """
  del vocab_size  # unused arg
  inputs = x
  p = model_hparams

  num_mel_bins = p.audio_num_mel_bins
  num_channels = 3 if p.audio_add_delta_deltas else 1

  with tf.variable_scope("speech_recognition_modality"):
    if p.audio_preproc_in_bottom:
      # Compute filterbanks
      with tf.variable_scope("fbanks"):
        waveforms = tf.squeeze(inputs, [2, 3])
        mel_fbanks = common_audio.compute_mel_filterbank_features(
            waveforms,
            sample_rate=p.audio_sample_rate,
            dither=p.audio_dither,
            preemphasis=p.audio_preemphasis,
            frame_length=p.audio_frame_length,
            frame_step=p.audio_frame_step,
            lower_edge_hertz=p.audio_lower_edge_hertz,
            upper_edge_hertz=p.audio_upper_edge_hertz,
            num_mel_bins=p.audio_num_mel_bins,
            apply_mask=True)
        if p.audio_add_delta_deltas:
          mel_fbanks = common_audio.add_delta_deltas(mel_fbanks)
        x = tf.reshape(mel_fbanks,
                       common_layers.shape_list(mel_fbanks)[:2] +
                       [num_mel_bins, num_channels])

        nonpadding_mask = 1. - common_attention.embedding_to_padding(x)
        num_of_nonpadding_elements = tf.reduce_sum(
            nonpadding_mask) * num_mel_bins * num_channels

        # This replaces CMVN estimation on data
        var_epsilon = 1e-09
        mean = tf.reduce_sum(
            x, axis=[1], keepdims=True) / num_of_nonpadding_elements
        variance = (num_of_nonpadding_elements * mean**2. -
                    2. * mean * tf.reduce_sum(x, axis=[1], keepdims=True) +
                    tf.reduce_sum(x**2, axis=[1], keepdims=True)
                   ) / num_of_nonpadding_elements
        x = (x - mean) * tf.rsqrt(variance + var_epsilon) * tf.expand_dims(
            nonpadding_mask, -1)
    else:
      x = inputs

    # The convention is that the models are flattened along the spatial,
    # dimensions, thus the speech preprocessor treats frequencies and
    # channels as image colors (last axis)
    x.set_shape([None, None, num_mel_bins, num_channels])

    # TODO(chorowski): how to specify bottom's hparams and avoid hardcoding?
    x = tf.pad(x, [[0, 0], [0, 8], [0, 0], [0, 0]])
    for _ in range(2):
      x = tf.layers.conv2d(
          x, 128, (3, 3), (2, 2), use_bias=False)
      x = common_layers.layer_norm(x)
      x = tf.nn.relu(x)

    xshape = common_layers.shape_list(x)
    # apply a conv that will remove all frequencies and at the same time
    # project the output into desired hidden_size
    x = tf.pad(x, [[0, 0], [0, 2], [0, 0], [0, 0]])
    x = tf.layers.conv2d(x, p.hidden_size, (3, xshape[2]), use_bias=False)

    assert common_layers.shape_list(x)[2] == 1
    x = common_layers.layer_norm(x)
    x = tf.nn.relu(x)
  return x


def get_weights(model_hparams, vocab_size, hidden_dim=None):
  """Create or get concatenated embedding or softmax variable.

  Args:
    model_hparams: HParams, model hyperparmeters.
    vocab_size: int, vocabulary size.
    hidden_dim: dim of the variable. Defaults to _model_hparams' hidden_size

  Returns:
     a list of num_shards Tensors.
  """
  if hidden_dim is None:
    hidden_dim = model_hparams.hidden_size
  num_shards = model_hparams.symbol_modality_num_shards
  shards = []
  for i in range(num_shards):
    shard_size = (vocab_size // num_shards) + (
        1 if i < vocab_size % num_shards else 0)
    var_name = "weights_%d" % i
    shards.append(
        tf.get_variable(
            var_name, [shard_size, hidden_dim],
            initializer=tf.random_normal_initializer(0.0, hidden_dim**-0.5)))
  if num_shards == 1:
    ret = shards[0]
  else:
    ret = tf.concat(shards, 0)
  # Convert ret to tensor.
  if not tf.executing_eagerly():
    ret = common_layers.convert_gradient_to_tensor(ret)
  return ret


def _symbol_bottom_simple(x, model_hparams, vocab_size, name, reuse):
  """Bottom transformation for symbols."""
  with tf.variable_scope(name, reuse=reuse):
    # Ensure the inputs are 3-D
    if len(x.get_shape()) == 4:
      x = tf.squeeze(x, axis=3)
    while len(x.get_shape()) < 3:
      x = tf.expand_dims(x, axis=-1)

    var = get_weights(model_hparams, vocab_size)
    x = common_layers.dropout_no_scaling(
        x, 1.0 - model_hparams.symbol_dropout)
    ret = common_layers.gather(var, x)
    if model_hparams.multiply_embedding_mode == "sqrt_depth":
      ret *= model_hparams.hidden_size**0.5
    ret *= tf.expand_dims(
        common_layers.cast_like(tf.not_equal(x, 0), ret), -1)
    return ret


def symbol_bottom(x, model_hparams, vocab_size):
  if (model_hparams.shared_embedding_and_softmax_weights or
      model_hparams.get("shared_embedding")):
    return _symbol_bottom_simple(
        x, model_hparams, vocab_size, "shared", reuse=None)
  return _symbol_bottom_simple(
      x, model_hparams, vocab_size, "input_emb", reuse=None)


def symbol_targets_bottom(x, model_hparams, vocab_size):
  """Bottom transformation for target symbols."""
  if (model_hparams.shared_embedding_and_softmax_weights or
      model_hparams.get("shared_embedding")):
    try:
      return _symbol_bottom_simple(
          x, model_hparams, vocab_size, "shared", reuse=True)
    except ValueError:
      # perhaps there were no inputs, and this is a new variable.
      return _symbol_bottom_simple(
          x, model_hparams, vocab_size, "shared", reuse=None)
  else:
    return _symbol_bottom_simple(
        x, model_hparams, vocab_size, "target_emb", reuse=None)


def symbol_one_hot_bottom(x, model_hparams, vocab_size):
  del model_hparams  # unused arg
  return tf.one_hot(x, vocab_size)


def video_bottom(x, model_hparams, vocab_size):
  del model_hparams, vocab_size  # unused arg
  common_video.gif_summary("inputs", x, max_outputs=1)
  x = common_layers.standardize_images(x)
  return x


def video_targets_bottom(x, model_hparams, vocab_size):
  del model_hparams, vocab_size  # unused arg
  common_video.gif_summary("targets", x, max_outputs=1)
  x = common_layers.standardize_images(x)
  return x


def video_bitwise_bottom(x, model_hparams, vocab_size):
  """Bottom transformation for embedding video bitwise."""
  pixel_embedding_size = 64
  inputs = x
  with tf.variable_scope("video_modality_bitwise", reuse=tf.AUTO_REUSE):
    common_layers.summarize_video(inputs, "bottom")
    # Embed bitwise.
    assert vocab_size == 256
    embedded = discretization.int_to_bit_embed(inputs, 8,
                                               pixel_embedding_size)
    # Project.
    return tf.layers.dense(
        embedded,
        model_hparams.hidden_size,
        name="merge_pixel_embedded_frames")


def video_bitwise_targets_bottom(x, model_hparams, vocab_size):
  """Bottom transformation for embedding target video bitwise."""
  pixel_embedding_size = 64
  inputs = x
  with tf.variable_scope("video_modality_bitwise", reuse=tf.AUTO_REUSE):
    common_layers.summarize_video(inputs, "targets_bottom")
    # Embed bitwise.
    assert vocab_size == 256
    embedded = discretization.int_to_bit_embed(inputs, 8,
                                               pixel_embedding_size)
    # Transpose and project.
    transposed = common_layers.time_to_channels(embedded)
    return tf.layers.dense(
        transposed,
        model_hparams.hidden_size,
        name="merge_pixel_embedded_frames")


def video_identity_bottom(x, model_hparams, vocab_size):
  del model_hparams, vocab_size  # unused arg
  common_video.gif_summary("inputs", x, max_outputs=1)
  return x


def video_identity_targets_bottom(x, model_hparams, vocab_size):
  del model_hparams, vocab_size  # unused arg
  common_video.gif_summary("targets", x, max_outputs=1)
  return x


def video_pixel_noise_bottom(x, model_hparams, vocab_size):
  """Bottom transformation for video."""
  input_noise = getattr(model_hparams, "video_modality_input_noise", 0.25)
  inputs = x
  if model_hparams.mode == tf.estimator.ModeKeys.TRAIN:
    background = tfp.stats.percentile(inputs, 50., axis=[0, 1, 2, 3])
    input_shape = common_layers.shape_list(inputs)
    input_size = tf.reduce_prod(input_shape[:-1])
    input_mask = tf.multinomial(
        tf.log([[input_noise, 1.-input_noise]]), input_size)
    input_mask = tf.reshape(tf.cast(input_mask, tf.int32),
                            input_shape[:-1]+[1])
    inputs = inputs * input_mask + background * (1 - input_mask)
  return video_bottom(inputs, model_hparams, vocab_size)


def convert_rgb_to_real(prediction, targets):
  """Convert prediction and target from rgb to real."""
  prediction = tf.squeeze(prediction, axis=-1)
  prediction = common_layers.convert_rgb_to_real(prediction)
  targets = common_layers.convert_rgb_to_real(targets)
  return prediction, targets


def video_raw_bottom(x, model_hparams, vocab_size):
  del model_hparams, vocab_size  # unused arg
  common_video.gif_summary("inputs", x)
  return common_layers.convert_rgb_to_real(x)


def video_raw_targets_bottom(x, model_hparams, vocab_size):
  del model_hparams, vocab_size  # unused arg
  common_video.gif_summary("targets_bottom", x)
  return common_layers.convert_rgb_to_real(x)


# Loss transformations, applied to target features


def ctc_symbol_loss(top_out, targets, model_hparams, vocab_size, weight_fn):
  """Compute the CTC loss."""
  del model_hparams, vocab_size  # unused arg
  logits = top_out
  with tf.name_scope("ctc_loss", values=[logits, targets]):
    # For CTC we assume targets are 1d, [batch, length, 1, 1] here.
    targets_shape = targets.get_shape().as_list()
    assert len(targets_shape) == 4
    assert targets_shape[2] == 1
    assert targets_shape[3] == 1
    targets = tf.squeeze(targets, axis=[2, 3])
    logits = tf.squeeze(logits, axis=[2, 3])
    targets_mask = 1 - tf.to_int32(tf.equal(targets, 0))
    targets_lengths = tf.reduce_sum(targets_mask, axis=1)
    sparse_targets = tf.keras.backend.ctc_label_dense_to_sparse(
        targets, targets_lengths)
    xent = tf.nn.ctc_loss(
        sparse_targets,
        logits,
        targets_lengths,
        time_major=False,
        preprocess_collapse_repeated=False,
        ctc_merge_repeated=False)
    weights = weight_fn(targets)
    return tf.reduce_sum(xent), tf.reduce_sum(weights)


def generic_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
  """Compute loss numerator and denominator for one shard of output."""
  del vocab_size  # unused arg
  logits = top_out
  logits = common_attention.maybe_upcast(logits, hparams=model_hparams)
  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.0)
  return common_layers.padded_cross_entropy(
      logits,
      targets,
      model_hparams.label_smoothing,
      cutoff=cutoff,
      weights_fn=weights_fn)


def generic_l2_loss(body_output,
                    targets,
                    model_hparams,
                    vocab_size,
                    weights_fn):
  del model_hparams, vocab_size, weights_fn  # unused arg
  loss = tf.squared_difference(body_output, tf.to_float(targets))
  return tf.reduce_mean(loss), tf.constant(1.0)


def multi_label_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
  """Average loss over the labels."""
  del vocab_size  # unused arg
  logits = top_out
  num_labels = tf.shape(targets)[1]
  logits = tf.tile(logits, [1, num_labels, 1, 1, 1])

  xent, weights = common_layers.padded_cross_entropy(
      logits,
      targets,
      model_hparams.label_smoothing,
      weights_fn=weights_fn,
      reduce_sum=False,
  )
  xent = tf.squeeze(xent, [2, 3])
  weights = tf.squeeze(weights, [2, 3])
  # average loss over all labels
  loss = tf.reduce_sum(xent, axis=1)
  weights = tf.reduce_sum(weights, axis=1)
  loss /= (weights + 1e-8)
  weights = tf.to_float(tf.greater(weights, 0.))

  return tf.reduce_sum(loss*weights), tf.reduce_sum(weights)


def one_hot_class_label_loss(top_out,
                             targets,
                             model_hparams,
                             vocab_size,
                             weights_fn):
  """Apply softmax cross-entropy between outputs and targets.

  Args:
    top_out: logits Tensor with shape [batch, ?, ?, num_classes]
    targets: one-hot encoding Tensor with shape [batch, ?, ?, num_classes]
    model_hparams: HParams, model hyperparmeters.
    vocab_size: int, vocabulary size.
    weights_fn:

  Returns:
    loss_scale (cross-entropy), loss_denom
  """
  del model_hparams, vocab_size  # unused arg
  loss_scale = tf.losses.softmax_cross_entropy(
      onehot_labels=targets, logits=top_out)
  weights = weights_fn(targets)
  loss_denom = tf.reduce_sum(weights)
  return loss_scale, loss_denom


def real_l2_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
  del model_hparams, vocab_size  # unused arg
  predictions = top_out
  if (len(common_layers.shape_list(top_out)) != len(
      common_layers.shape_list(targets))):
    predictions = tf.squeeze(top_out, axis=[-1])
  with tf.name_scope("l2"):
    weights = weights_fn(targets)
    l2 = tf.pow(predictions - targets, 2)
    return tf.reduce_sum(l2 * weights), tf.reduce_sum(weights)


def real_log_poisson_loss(top_out,
                          targets,
                          model_hparams,
                          vocab_size,
                          weights_fn):
  """Poisson loss for real."""
  del model_hparams, vocab_size  # unused arg
  predictions = top_out
  if (len(common_layers.shape_list(top_out)) != len(
      common_layers.shape_list(targets))):
    predictions = tf.squeeze(top_out, axis=[-1])
  with tf.name_scope("log_possion"):
    weights = weights_fn(targets)
    lp_loss = tf.nn.log_poisson_loss(targets, predictions)
    return tf.reduce_sum(lp_loss * weights), tf.reduce_sum(weights)


def sigmoid_class_label_loss(top_out,
                             targets,
                             model_hparams,
                             vocab_size,
                             weights_fn):
  """Loss for class label."""
  # Expect inputs of size [batch-size, timesteps, 1, num-classes], where the
  # last dimension of num-classes represents logits for binary labels
  del model_hparams, vocab_size  # unused arg
  loss_scale = tf.losses.sigmoid_cross_entropy(
      multi_class_labels=targets, logits=top_out)
  weights = weights_fn(targets)
  loss_denom = tf.reduce_sum(weights)
  return loss_scale, loss_denom


def sigmoid_max_pooling_class_label_loss(top_out,
                                         targets,
                                         model_hparams,
                                         vocab_size,
                                         weights_fn):
  """Loss for class label."""
  # Expect inputs of size [batch-size, 1, 1, num-classes], where the
  # last dimension of num-classes represents logits for binary labels
  del model_hparams, vocab_size  # unused arg
  loss_scale = tf.losses.sigmoid_cross_entropy(
      multi_class_labels=targets, logits=top_out)
  weights = weights_fn(targets)
  loss_denom = tf.reduce_sum(weights)
  return loss_scale, loss_denom


def symbol_one_hot_loss(top_out,
                        targets,
                        model_hparams,
                        vocab_size,
                        weights_fn):
  del model_hparams, weights_fn  # unused arg
  labels = tf.one_hot(targets, vocab_size)
  loss = tf.nn.softmax_cross_entropy_with_logits(
      logits=top_out, labels=labels)
  return tf.reduce_mean(loss), tf.constant(1.0)


def video_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
  """Compute loss numerator and denominator for one shard of output."""
  del vocab_size  # unused arg
  logits = top_out
  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
  return common_layers.padded_cross_entropy(
      logits,
      targets,
      model_hparams.label_smoothing,
      cutoff=cutoff,
      weights_fn=weights_fn)


def video_identity_loss(top_out,
                        targets,
                        model_hparams,
                        vocab_size,
                        weights_fn):
  """Compute loss numerator and denominator for one shard of output."""
  del vocab_size  # unused arg
  # TODO(nikip): Try L2 loss
  logits = top_out
  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:])
  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.01)
  return common_layers.padded_cross_entropy(
      logits,
      targets,
      model_hparams.label_smoothing,
      cutoff=cutoff,
      weights_fn=weights_fn)


def video_l1_internal_loss(logits, targets, model_hparams):
  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
  return tf.nn.relu(tf.abs(logits - targets) - cutoff)


def video_l1_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
  """Compute loss numerator and denominator for one shard of output."""
  del vocab_size  # unused arg
  logits = top_out
  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
  weights = weights_fn(targets)
  # Shift targets by 0.5 so later just casting to int gives the prediction.
  # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
  # Later (in merics or infer) this is cast to int anyway. Also, we have no
  # loss beyond cutoff = 0.2 as these are already correct predictions.
  targets = tf.to_float(targets) + 0.5
  loss = video_l1_internal_loss(logits, targets, model_hparams)
  return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)


def video_l2_internal_loss(logits, targets, model_hparams):
  cutoff = getattr(model_hparams, "video_modality_loss_cutoff", 0.2)
  return tf.nn.relu(
      tf.squared_difference(logits, targets) - cutoff * cutoff)


def video_l2_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
  """Compute loss numerator and denominator for one shard of output."""
  del vocab_size  # unused arg
  logits = top_out
  logits = tf.reshape(logits, [-1] + common_layers.shape_list(logits)[2:-1])
  targets = tf.reshape(targets, [-1] + common_layers.shape_list(targets)[2:])
  weights = weights_fn(targets)
  # Shift targets by 0.5 so later just casting to int gives the prediction.
  # So for int targets, say 0 and 7, we actually train to predict 0.5 and 7.5.
  # Later (in merics or infer) this is cast to int anyway. Also, we have no
  # loss beyond cutoff = 0.2 as these are already correct predictions.
  targets = tf.to_float(targets) + 0.5
  loss = video_l2_internal_loss(logits, targets, model_hparams)
  return tf.reduce_sum(loss * weights), tf.reduce_sum(weights)


def video_l2_raw_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
  del model_hparams, vocab_size, weights_fn  # unused arg
  prediction, groundtruth = convert_rgb_to_real(top_out, targets)
  loss = tf.losses.mean_squared_error(prediction, groundtruth)
  return loss, tf.constant(1.0)


def video_l1_raw_loss(top_out, targets, model_hparams, vocab_size, weights_fn):
  del model_hparams, vocab_size, weights_fn  # unused arg
  prediction, groundtruth = convert_rgb_to_real(top_out, targets)
  loss = tf.losses.absolute_difference(prediction, groundtruth)
  return loss, tf.constant(1.0)


# Top transformations, applied to target features


def is_pointwise(func):
  """Decorator for whether the function is pointwise.

  An example of a pointwise function is a linear layer followed by
  a softmax. Given a tensor [batch, length, height, depth] it operates
  only on the last axis, on every point in [batch, length, height] fully
  independently. In contrast, a classifier that first averages over length
  and height is not pointwise, as it depends on the whole field. It is useful
  to know if top functions are pointwise to speed up decoding in certain models.

  Args:
    func: Function to decorate.

  Returns:
    Original function with an attribute pointwise set to True.
  """
  func.pointwise = True
  return func


def class_label_top(body_output, targets, model_hparams, vocab_size):
  """Transform inputs from model space to target space.

  Average over inner dims and a linear layer to logits.

  Args:
    body_output: A Tensor with shape [batch, ?, ?, body_output_size].
    targets:
    model_hparams: HParams, model hyperparmeters.
    vocab_size: int, vocabulary size.

  Returns:
    a Tensors, each with shape [batch_size, 1, 1, 1, vocab_size]
  """
  del targets  # unused arg
  with tf.variable_scope("class_label_modality_%d_%d" % (
      vocab_size, model_hparams.hidden_size)):
    x = body_output
    x = tf.reduce_mean(x, axis=[1, 2], keepdims=True)
    res = tf.layers.dense(x, vocab_size)
    return tf.expand_dims(res, 3)


def identity_top(body_output, targets, model_hparams, vocab_size):
  del targets, model_hparams, vocab_size  # unused arg
  return body_output


def image_top(body_output, targets, model_hparams, vocab_size):
  """Top transformation for images."""
  del targets  # unused arg
  # TODO(lukaszkaiser): is this a universal enough way to get channels?
  num_channels = model_hparams.problem.num_channels
  with tf.variable_scope("rgb_softmax"):
    body_output_shape = common_layers.shape_list(body_output)
    reshape_shape = body_output_shape[:3]
    reshape_shape.extend([num_channels, vocab_size])
    res = tf.layers.dense(body_output, vocab_size * num_channels)
    res = tf.reshape(res, reshape_shape)
    if not tf.get_variable_scope().reuse:
      res_argmax = tf.argmax(res, axis=-1)
      tf.summary.image(
          "result",
          common_layers.tpu_safe_image_summary(res_argmax),
          max_outputs=1)
    return res


def image_channel_compress_top(body_output, targets, model_hparams, vocab_size):
  """Transforms body output to return logits.

  Args:
    body_output: Tensor of shape [batch, img_len, img_len, depth].
    targets:
    model_hparams: HParams, model hyperparmeters.
    vocab_size: int, vocabulary size.

  Returns:
    Tensor of shape [batch, img_len, img_len, channels, vocab_size].
  """
  del targets  # unused arg
  with tf.variable_scope("image_channel_compress_modality"):
    hidden_size = model_hparams.hidden_size
    img_len = model_hparams.img_len
    channels = 3  # RGB
    batch = common_layers.shape_list(body_output)[0]
    x = tf.layers.conv2d(
        body_output,
        hidden_size * channels,
        kernel_size=(1, 1),
        strides=(1, 1),
        padding="VALID",
        activation=tf.nn.relu,
        name="decompress_conv")
    x = tf.reshape(x, [batch, img_len, img_len * channels, hidden_size])
    x = common_layers.layer_preprocess(x, model_hparams)
    x = tf.layers.dense(x,
                        vocab_size,
                        use_bias=True,
                        activation=None,
                        name="output_conv")
    x = tf.reshape(
        x, [batch, img_len, img_len, channels, vocab_size])
    return x


def image_channel_embeddings_top(body_output,
                                 targets,
                                 model_hparams,
                                 vocab_size):
  """Top transformation for images."""
  del targets  # unused arg
  with tf.variable_scope("image_channel_embeddings_bottom"):
    img_len = model_hparams.img_len
    channels = model_hparams.num_channels
    x = tf.layers.dense(
        body_output, 256, use_bias=True, activation=None, name="output_conv")
    x = tf.reshape(x,
                   [-1, img_len, img_len, channels, vocab_size])
    return x


@is_pointwise
def real_top(body_output, targets, model_hparams, vocab_size):
  del targets, model_hparams  # unused arg
  with tf.variable_scope("real"):
    return tf.layers.dense(body_output, vocab_size, name="top")


def sigmoid_max_pooling_class_label_top(body_output,
                                        targets,
                                        model_hparams,
                                        vocab_size):
  """Transform inputs from model space to target space.

  Average over inner dims and a linear layer to logits.

  Args:
    body_output: A Tensor with shape [batch, timesteps, 1, body_output_size].
    targets:
    model_hparams: HParams, model hyperparmeters.
    vocab_size: int, vocabulary size.

  Returns:
    a Tensors, each with shape [batch_size, 1, 1, vocab_size]
  """
  del targets  # unused arg
  with tf.variable_scope(
      "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
          vocab_size, model_hparams.hidden_size)):
    x = body_output
    x = tf.reduce_max(x, axis=1, keepdims=True)
    return tf.layers.dense(x, vocab_size)


def softmax_average_pooling_class_label_top(body_output,
                                            targets,
                                            model_hparams,
                                            vocab_size):
  """Loss for class label."""
  del targets  # unused arg
  with tf.variable_scope(
      "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
          vocab_size, model_hparams.hidden_size)):
    x = body_output
    x = tf.reduce_mean(x, axis=1, keepdims=True)
    return tf.layers.dense(x, vocab_size)


def softmax_last_timestep_class_label_top(body_output,
                                          targets,
                                          model_hparams,
                                          vocab_size):
  """Loss for class label."""
  del targets  # unused arg
  with tf.variable_scope(
      "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
          vocab_size, model_hparams.hidden_size)):
    x = body_output
    x = tf.expand_dims(x[:, -1], 1)  # Pick the last timestep
    return tf.layers.dense(x, vocab_size)


def softmax_max_pooling_class_label_top(body_output,
                                        targets,
                                        model_hparams,
                                        vocab_size):
  """Loss for class label."""
  del targets  # unused arg
  with tf.variable_scope(
      "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
          vocab_size, model_hparams.hidden_size)):
    x = body_output
    x = tf.reduce_max(x, axis=1, keepdims=True)
    return tf.layers.dense(x, vocab_size)


@is_pointwise
def symbol_top(body_output, targets, model_hparams, vocab_size):
  """Generate logits.

  Args:
    body_output: A Tensor with shape
      [batch, p0, p1, model_hparams.hidden_size].
    targets: Unused.
    model_hparams: HParams, model hyperparmeters.
    vocab_size: int, vocabulary size.

  Returns:
    logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
  """
  del targets  # unused arg
  if model_hparams.shared_embedding_and_softmax_weights:
    scope_name = "shared"
    reuse = tf.AUTO_REUSE
  else:
    scope_name = "softmax"
    reuse = False
  with tf.variable_scope(scope_name, reuse=reuse):
    body_output_shape = common_layers.shape_list(body_output)
    var = get_weights(model_hparams, vocab_size, body_output_shape[-1])
    if (model_hparams.factored_logits and
        model_hparams.mode == tf.estimator.ModeKeys.TRAIN):
      # insert channels dimension
      body_output = tf.expand_dims(body_output, 3)
      return common_layers.FactoredTensor(body_output, var)
    else:
      body_output = tf.reshape(body_output, [-1, body_output_shape[-1]])
      logits = tf.matmul(body_output, var, transpose_b=True)
      return tf.reshape(logits,
                        body_output_shape[:-1] + [1, vocab_size])


@is_pointwise
def symbol_one_hot_top(body_output, targets, model_hparams, vocab_size):
  del targets, model_hparams, vocab_size  # unused arg
  return body_output


def video_top(body_output, targets, model_hparams, vocab_size):
  """Top transformation for video."""
  del targets  # unused arg
  num_channels = model_hparams.problem.num_channels
  shape = common_layers.shape_list(body_output)
  reshape_shape = shape[:-1] + [num_channels, vocab_size]
  res = tf.reshape(body_output, reshape_shape)
  # Calculate argmax so as to have a summary with the produced images.
  x = tf.argmax(tf.reshape(res, [-1, vocab_size]), axis=-1)
  x = tf.reshape(x, shape[:-1] + [num_channels])
  common_video.gif_summary("results", x, max_outputs=1)
  return res


def video_l1_top(body_output, targets, model_hparams, vocab_size):
  """Top transformation for video."""
  del targets, vocab_size  # unused arg
  num_channels = model_hparams.problem.num_channels
  num_frames = model_hparams.video_num_target_frames
  with tf.variable_scope("rgb"):
    body_output_shape = common_layers.shape_list(body_output)
    res = tf.layers.dense(body_output, num_channels * num_frames, name="cast")
    res = tf.reshape(res, body_output_shape[:3] + [num_channels, num_frames])
    res = tf.transpose(res, [0, 4, 1, 2, 3])  # Move frames next to batch.
    if not tf.get_variable_scope().reuse:
      res_argmax = res[:, -1, :, :, :]
      tf.summary.image(
          "result",
          common_layers.tpu_safe_image_summary(res_argmax),
          max_outputs=1)
    return tf.expand_dims(res, axis=-1)  # Add an axis like in perplexity.


def video_raw_top(body_output, targets, model_hparams, vocab_size):
  del targets, model_hparams, vocab_size  # unused arg
  frames = body_output
  if isinstance(body_output, list):
    frames = tf.stack(body_output, axis=1)
  rgb_frames = common_layers.convert_real_to_rgb(frames)
  common_video.gif_summary("body_output", rgb_frames)
  return tf.expand_dims(rgb_frames, axis=-1)


# Utility functions similar to tf.keras for default transformations


def get_bottom(modality_type, value=None):
  """Gets default bottom transformation; if none available, return value."""
  if modality_type == ModalityType.AUDIO:
    return audio_bottom
  elif modality_type == ModalityType.AUDIO_SPECTRAL:
    return audio_spectral_bottom
  elif modality_type in (ModalityType.CLASS_LABEL,
                         ModalityType.MULTI_LABEL,
                         ModalityType.ONE_HOT_CLASS_LABEL,
                         ModalityType.SIGMOID_CLASS_LABEL,
                         ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
    return class_label_bottom
  elif modality_type in (ModalityType.CTC_SYMBOL,
                         ModalityType.SYMBOL,
                         ModalityType.SYMBOL_WEIGHTS_ALL):
    return symbol_bottom
  elif modality_type in (ModalityType.GENERIC_L2_LOSS,
                         ModalityType.IDENTITY,
                         ModalityType.IDENTITY_SYMBOL,
                         ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM):
    return identity_bottom
  elif modality_type == ModalityType.IMAGE:
    return image_bottom
  elif modality_type in (ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
                         ModalityType.IMAGE_CHANNEL_COMPRESS):
    return image_channel_compress_bottom
  elif modality_type in (ModalityType.REAL,
                         ModalityType.REAL_L2_LOSS,
                         ModalityType.REAL_LOG_POISSON_LOSS):
    return real_bottom
  elif modality_type == ModalityType.SPEECH_RECOGNITION:
    return speech_recognition_bottom
  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
    return symbol_one_hot_bottom
  elif modality_type in (ModalityType.VIDEO,
                         ModalityType.VIDEO_L1,
                         ModalityType.VIDEO_L2):
    return video_bottom
  elif modality_type == ModalityType.VIDEO_BITWISE:
    return video_bitwise_bottom
  elif modality_type == ModalityType.VIDEO_IDENTITY:
    return video_identity_bottom
  elif modality_type in (ModalityType.VIDEO_L1_RAW,
                         ModalityType.VIDEO_L2_RAW):
    return video_raw_bottom
  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
    return video_pixel_noise_bottom
  return value


def get_loss(modality_type, value=None):
  """Gets default loss transformation; if none available, return value."""
  if modality_type in (ModalityType.AUDIO,
                       ModalityType.AUDIO_SPECTRAL,
                       ModalityType.CLASS_LABEL,
                       ModalityType.IDENTITY,
                       ModalityType.IDENTITY_SYMBOL,
                       ModalityType.IMAGE,
                       ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
                       ModalityType.IMAGE_CHANNEL_COMPRESS,
                       ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM,
                       ModalityType.REAL,
                       ModalityType.SPEECH_RECOGNITION,
                       ModalityType.SYMBOL,
                       ModalityType.SYMBOL_WEIGHTS_ALL):
    return generic_loss
  elif modality_type == ModalityType.CTC_SYMBOL:
    return ctc_symbol_loss
  elif modality_type == ModalityType.GENERIC_L2_LOSS:
    return generic_l2_loss
  elif modality_type == ModalityType.MULTI_LABEL:
    return multi_label_loss
  elif modality_type in (ModalityType.ONE_HOT_CLASS_LABEL,
                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
    return one_hot_class_label_loss
  elif modality_type == ModalityType.REAL_L2_LOSS:
    return real_l2_loss
  elif modality_type == ModalityType.REAL_LOG_POISSON_LOSS:
    return real_log_poisson_loss
  elif modality_type == ModalityType.SIGMOID_CLASS_LABEL:
    return sigmoid_class_label_loss
  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
    return sigmoid_max_pooling_class_label_loss
  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
    return symbol_one_hot_loss
  elif modality_type in (ModalityType.VIDEO,
                         ModalityType.VIDEO_BITWISE,
                         ModalityType.VIDEO_PIXEL_NOISE):
    return video_loss
  elif modality_type == ModalityType.VIDEO_IDENTITY:
    return video_identity_loss
  elif modality_type == ModalityType.VIDEO_L1:
    return video_l1_loss
  elif modality_type == ModalityType.VIDEO_L1_RAW:
    return video_l1_raw_loss
  elif modality_type == ModalityType.VIDEO_L2:
    return video_l2_loss
  elif modality_type == ModalityType.VIDEO_L2_RAW:
    return video_l2_raw_loss
  return value


def get_name(modality_type, value=None):
  """Gets default name for transformations; if none available, return value."""
  # For legacy reasons, modalities vary in their naming scheme. Future plans are
  # to remove any need for get_name. We do not recommend using it.
  if modality_type == ModalityType.AUDIO:
    return lambda model_hparams, vocab_size: "audio_modality"
  elif modality_type == ModalityType.AUDIO_SPECTRAL:
    return lambda model_hparams, vocab_size: "audio_spectral_modality"
  elif modality_type == ModalityType.GENERIC_L2_LOSS:
    return lambda model_hparams, vocab_size: "generic_l2_loss_modality"
  elif modality_type == ModalityType.IDENTITY:
    return lambda model_hparams, vocab_size: "identity_modality"
  elif modality_type == ModalityType.IMAGE:
    return lambda model_hparams, vocab_size: "image_modality"
  elif modality_type == ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY:
    return (lambda model_hparams, vocab_size:  # pylint: disable=g-long-lambda
            "image_channel_bottom_identity_modality")
  elif modality_type == ModalityType.IMAGE_CHANNEL_COMPRESS:
    return lambda model_hparams, vocab_size: "image_channel_compress_modality"
  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
    return lambda model_hparams, vocab_size: "image_channel_embeddings_bottom"
  elif modality_type == ModalityType.REAL:
    return lambda model_hparams, vocab_size: "real_modality"
  elif modality_type == ModalityType.REAL_L2_LOSS:
    return lambda model_hparams, vocab_size: "real_l2_loss_modality"
  elif modality_type == ModalityType.REAL_LOG_POISSON_LOSS:
    return lambda model_hparams, vocab_size: "real_log_poisson_loss_modality"
  elif modality_type == ModalityType.SPEECH_RECOGNITION:
    return lambda model_hparams, vocab_size: "speech_recognition_modality"
  elif modality_type == ModalityType.VIDEO:
    return lambda model_hparams, vocab_size: "video_modality"
  elif modality_type == ModalityType.VIDEO_BITWISE:
    return lambda model_hparams, vocab_size: "video_modality_bitwise"
  elif modality_type == ModalityType.VIDEO_IDENTITY:
    return lambda model_hparams, vocab_size: "video_modality_identity"
  elif modality_type == ModalityType.VIDEO_L1:
    return lambda model_hparams, vocab_size: "video_modality_l1"
  elif modality_type == ModalityType.VIDEO_L1_RAW:
    return lambda model_hparams, vocab_size: "video_modality_l1_raw"
  elif modality_type == ModalityType.VIDEO_L2:
    return lambda model_hparams, vocab_size: "video_modality_l2"
  elif modality_type == ModalityType.VIDEO_L2_RAW:
    return lambda model_hparams, vocab_size: "video_modality_l2_raw"
  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
    return lambda model_hparams, vocab_size: "video_modality_pixel_noise"
  elif modality_type in (ModalityType.CLASS_LABEL,
                         ModalityType.MULTI_LABEL,
                         ModalityType.ONE_HOT_CLASS_LABEL):
    def name(model_hparams, vocab_size):
      return "class_label_modality_%d_%d" % (vocab_size,
                                             model_hparams.hidden_size)
    return name
  elif modality_type in (ModalityType.CTC_SYMBOL,
                         ModalityType.IDENTITY_SYMBOL,
                         ModalityType.SYMBOL,
                         ModalityType.SYMBOL_WEIGHTS_ALL,
                         ModalityType.SYMBOL_ONE_HOT):
    def name(model_hparams, vocab_size):
      return "symbol_modality_%d_%d" % (vocab_size, model_hparams.hidden_size)
    return name
  elif modality_type == ModalityType.SIGMOID_CLASS_LABEL:
    def name(model_hparams, vocab_size):
      return "sigmoid_class_symbol_modality_%d_%d" % (vocab_size,
                                                      model_hparams.hidden_size)
    return name
  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
    def name(model_hparams, vocab_size):
      return "sigmoid_max_pooling_class_symbol_modality_%d_%d" % (
          vocab_size, model_hparams.hidden_size)
    return name
  elif modality_type == ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL:
    def name(model_hparams, vocab_size):
      return "softmax_average_pooling_onehot_class_label_modality_%d_%d" % (
          vocab_size, model_hparams.hidden_size)
    return name
  elif modality_type == ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL:
    def name(model_hparams, vocab_size):
      return "softmax_last_timestep_onehot_class_label_modality_%d_%d" % (
          vocab_size, model_hparams.hidden_size)
    return name
  elif modality_type == ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL:
    def name(model_hparams, vocab_size):
      return "softmax_max_pooling_onehot_class_label_modality_%d_%d" % (
          vocab_size, model_hparams.hidden_size)
    return name
  return value


def get_targets_bottom(modality_type, value=None):
  """Gets default bottom transformation for targets; if none, return value."""
  if modality_type == ModalityType.AUDIO:
    return make_targets_bottom(audio_bottom)
  elif modality_type == ModalityType.AUDIO_SPECTRAL:
    return make_targets_bottom(audio_spectral_bottom)
  elif modality_type in (ModalityType.CLASS_LABEL,
                         ModalityType.MULTI_LABEL,
                         ModalityType.ONE_HOT_CLASS_LABEL,
                         ModalityType.SIGMOID_CLASS_LABEL,
                         ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL,
                         ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL,
                         ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL,
                         ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL):
    return class_label_targets_bottom
  elif modality_type in (ModalityType.CTC_SYMBOL,
                         ModalityType.SYMBOL,
                         ModalityType.SYMBOL_WEIGHTS_ALL):
    return symbol_targets_bottom
  elif modality_type in (ModalityType.GENERIC_L2_LOSS,
                         ModalityType.IDENTITY_SYMBOL):
    return identity_bottom
  elif modality_type == ModalityType.IDENTITY:
    return make_targets_bottom(identity_bottom)
  elif modality_type == ModalityType.IMAGE:
    return image_targets_bottom
  elif modality_type in (ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
                         ModalityType.IMAGE_CHANNEL_COMPRESS):
    return image_channel_compress_targets_bottom
  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
    return image_channel_embeddings_bottom
  elif modality_type in (ModalityType.REAL,
                         ModalityType.REAL_L2_LOSS,
                         ModalityType.REAL_LOG_POISSON_LOSS):
    return make_targets_bottom(real_bottom)
  elif modality_type == ModalityType.SPEECH_RECOGNITION:
    return make_targets_bottom(speech_recognition_bottom)
  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
    return symbol_one_hot_bottom
  elif modality_type in (ModalityType.VIDEO,
                         ModalityType.VIDEO_L1,
                         ModalityType.VIDEO_L2):
    return video_targets_bottom
  elif modality_type == ModalityType.VIDEO_BITWISE:
    return video_bitwise_targets_bottom
  elif modality_type == ModalityType.VIDEO_IDENTITY:
    return video_identity_targets_bottom
  elif modality_type in (ModalityType.VIDEO_L1_RAW,
                         ModalityType.VIDEO_L2_RAW):
    return video_raw_targets_bottom
  elif modality_type == ModalityType.VIDEO_PIXEL_NOISE:
    return make_targets_bottom(video_pixel_noise_bottom)
  return value


def get_top(modality_type, value=None):
  """Gets default top transformation; if none available, return value."""
  if modality_type in (ModalityType.AUDIO,
                       ModalityType.AUDIO_SPECTRAL,
                       ModalityType.GENERIC_L2_LOSS,
                       ModalityType.IDENTITY,
                       ModalityType.IDENTITY_SYMBOL,
                       ModalityType.IMAGE_CHANNEL_BOTTOM_IDENTITY,
                       ModalityType.SPEECH_RECOGNITION,
                       ModalityType.VIDEO_IDENTITY):
    return identity_top
  elif modality_type in (ModalityType.CLASS_LABEL,
                         ModalityType.MULTI_LABEL,
                         ModalityType.ONE_HOT_CLASS_LABEL,
                         ModalityType.SIGMOID_CLASS_LABEL):
    return class_label_top
  elif modality_type in (ModalityType.CTC_SYMBOL,
                         ModalityType.SYMBOL,
                         ModalityType.SYMBOL_WEIGHTS_ALL):
    return symbol_top
  elif modality_type == ModalityType.IMAGE:
    return image_top
  elif modality_type == ModalityType.IMAGE_CHANNEL_COMPRESS:
    return image_channel_compress_top
  elif modality_type == ModalityType.IMAGE_CHANNEL_EMBEDDINGS_BOTTOM:
    return image_channel_embeddings_top
  elif modality_type in (ModalityType.REAL,
                         ModalityType.REAL_L2_LOSS,
                         ModalityType.REAL_LOG_POISSON_LOSS):
    return real_top
  elif modality_type == ModalityType.SIGMOID_MAX_POOLING_CLASS_LABEL:
    return sigmoid_max_pooling_class_label_top
  elif modality_type == ModalityType.SOFTMAX_AVERAGE_POOLING_CLASS_LABEL:
    return softmax_average_pooling_class_label_top
  elif modality_type == ModalityType.SOFTMAX_LAST_TIMESTEP_CLASS_LABEL:
    return softmax_last_timestep_class_label_top
  elif modality_type == ModalityType.SOFTMAX_MAX_POOLING_CLASS_LABEL:
    return softmax_max_pooling_class_label_top
  elif modality_type == ModalityType.SYMBOL_ONE_HOT:
    return symbol_one_hot_top
  elif modality_type in (ModalityType.VIDEO,
                         ModalityType.VIDEO_BITWISE,
                         ModalityType.VIDEO_PIXEL_NOISE):
    return video_top
  elif modality_type in (ModalityType.VIDEO_L1,
                         ModalityType.VIDEO_L2):
    return video_l1_top
  elif modality_type in (ModalityType.VIDEO_L1_RAW,
                         ModalityType.VIDEO_L2_RAW):
    return video_raw_top
  return value


def get_weights_fn(modality_type, value=None):
  """Gets default weights function; if none available, return value."""
  if modality_type in (ModalityType.CTC_SYMBOL,
                       ModalityType.IDENTITY_SYMBOL,
                       ModalityType.MULTI_LABEL,
                       ModalityType.SYMBOL,
                       ModalityType.SYMBOL_ONE_HOT):
    return common_layers.weights_nonzero
  elif modality_type in ModalityType.get_choices():
    return common_layers.weights_all
  return value
