# coding=utf-8
# Copyright 2022 The Google Research 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.

"""Utilities for Tensorflow 2.0.

Partially adapted from:
https://www.tensorflow.org/tutorials/text/image_captioning
"""
# Lint as: python3
# pylint: disable=invalid-name

from __future__ import absolute_import
from __future__ import division

import tensorflow as tf


def film_params(sentence_embedding, n_layer_channel):
  """Generate FiLM parameters from a sentence embedding.

  Generate FiLM parameters from a sentence embedding. This method assumes a
  batch dimension exists.

  Args:
    sentence_embedding: a tensor containing batched sentenced embedding to be
      transformed
    n_layer_channel:    a list of integers specifying how many channels are at
      each hidden layer to be FiLM'ed

  Returns:
    a tuple of tensors the same length as n_layer_channel. Each element
    contains all gamma_i and beta_i for a single hidden layer.
  """
  n_total = sum(n_layer_channel) * 2
  all_params = tf.layers.dense(sentence_embedding, n_total)
  all_params = tf.keras.layers.Dense(
      2 * sum * (n_layer_channel), activation=tf.nn.relu)
  return tf.split(all_params, [c * 2 for c in n_layer_channel], 1)


def stack_conv_layer(layer_cfg, padding='same'):
  """Stack convolution layers per layer_cfg.

  Args:
    layer_cfg: list of integer tuples specifying the parameter each layer;
      each tuple should be (channel, kernel size, strides)
    padding: what kind of padding the conv layers use

  Returns:
    the keras model with stacked conv layers
  """
  layers = []
  for cfg in layer_cfg[:-1]:
    layers.append(
        tf.keras.layers.Conv2D(
            filters=cfg[0],
            kernel_size=cfg[1],
            strides=cfg[2],
            activation=tf.nn.relu,
            padding=padding))
  final_cfg = layer_cfg[-1]
  layers.append(
      tf.keras.layers.Conv2D(
          final_cfg[0], final_cfg[1], final_cfg[2], padding=padding))
  return tf.keras.Sequential(layers)


def stack_dense_layer(layer_cfg):
  """Stack Dense layers.

  Args:
    layer_cfg: list of integer specifying the number of units at each layer

  Returns:
    the keras model with stacked dense layers
  """
  layers = []
  for cfg in layer_cfg[:-1]:
    layers.append(tf.keras.layers.Dense(cfg, activation=tf.nn.relu))
  layers.append(tf.keras.layers.Dense(layer_cfg[-1]))
  return tf.keras.Sequential(layers)


def soft_variables_update(source_variables, target_variables, polyak_rate=1.0):
  """Update the target variables using exponential moving average.

  Specifically, v_s' = v_s * polyak_rate + (1-polyak_rate) * v_t

  Args:
    source_variables:  the moving average variables
    target_variables:  the new observations
    polyak_rate: rate of moving average

  Returns:
    Operation that does the update
  """
  updates = []
  for (v_s, v_t) in zip(source_variables, target_variables):
    v_t.shape.assert_is_compatible_with(v_s.shape)

    def update_fn(v1, v2):
      """Update variables."""
      # For not trainable variables do hard updates.
      return v1.assign(polyak_rate * v1 + (1 - polyak_rate) * v2)

    update = update_fn(v_t, v_s)
    updates.append(update)
  return updates


def vector_tensor_product(a, b):
  """"Returns keras layer that perfrom a outer product between a and b."""
  # a shape: [B, ?, d], b shape: [B, ?, d]
  shape_layer = tf.keras.layers.Lambda(tf.shape)
  shape = shape_layer(b)
  shape_numpy = b.get_shape()
  variable_length = shape[1]  # variable_len = ?
  expand_dims_layer_1 = tf.keras.layers.Reshape((-1, 1, shape_numpy[-1]))
  expand_dims_layer_2 = tf.keras.layers.Reshape((-1, 1, shape_numpy[-1]))
  a = expand_dims_layer_1(a)  # a shape: [B, ?, 1, d]
  b = expand_dims_layer_2(b)  # a shape: [B, ?, 1, d]
  tile_layer = tf.keras.layers.Lambda(
      lambda inputs: tf.tile(inputs[0], multiples=inputs[1]))
  a = tile_layer((a, [1, 1, variable_length, 1]))  # a shape: [B, ?, ?, d]
  b = tile_layer((b, [1, 1, variable_length, 1]))  # b shape: [B, ?, ?, d]
  b = tf.keras.layers.Permute((2, 1, 3))(b)  # b shape: [B, ?, ?, d]
  return tf.keras.layers.concatenate([a, b])  # shape: [B, ?, ?, 2*d]


class BahdanauAttention(tf.keras.Model):
  """Bahdanau Attention Layer.

  Attributes:
    w1: weights that process the feature
    w2: weights that process the memory state
    v: projection layer that project score vector to scalar
  """

  def __init__(self, units):
    """Initialize Bahdanau attention layer.

    Args:
      units: size of the dense layers
    """
    super(BahdanauAttention, self).__init__()
    self.W1 = tf.keras.layers.Dense(units)
    self.W2 = tf.keras.layers.Dense(units)
    self.V = tf.keras.layers.Dense(1)

  def call(self, features, hidden):
    # features(CNN_encoder output) shape == (batch_size, 64, embedding_dim)

    # hidden shape == (batch_size, hidden_size)
    # hidden_with_time_axis shape == (batch_size, 1, hidden_size)
    hidden_with_time_axis = tf.expand_dims(hidden, 1)

    # score shape == (batch_size, 64, hidden_size)
    score = tf.nn.tanh(self.W1(features) + self.W2(hidden_with_time_axis))

    # attention_weights shape == (batch_size, 64, 1)
    # you get 1 at the last axis because you are applying score to self.V
    attention_weights = tf.nn.softmax(self.V(score), axis=1)

    # context_vector shape after sum == (batch_size, hidden_size)
    context_vector = attention_weights * features
    context_vector = tf.reduce_sum(context_vector, axis=1)

    return context_vector, attention_weights


class GRUEnecoder(tf.keras.Model):
  """TF2.0 GRE encoder.

  Attributes:
    embedding: word embedding matrix
    gru: the GRU layer
  """

  def __init__(self, embedding_dim, units, vocab_size):
    """Initialize the GRU encoder.

    Args:
      embedding_dim: dimension of word emebdding
      units: number of units of the memory state
      vocab_size: total number of vocabulary
    """
    super(GRUEnecoder, self).__init__()
    self._units = units

    self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
    self.gru = tf.keras.layers.GRU(
        self.units,
        return_sequences=True,
        return_state=True,
        recurrent_initializer='glorot_uniform')

  def call(self, x, hidden):
    # x shape after passing through embedding == (batch_size, 1, embedding_dim)
    x = self.embedding(x)
    # passing the concatenated vector to the GRU
    output, state = self.gru(x)
    return output, state

  def reset_state(self, batch_size):
    return tf.zeros((batch_size, self._units))

[((47, 15, 47, 59), 'tensorflow.layers.dense', 'tf.layers.dense', ({(47, 31, 47, 49): 'sentence_embedding', (47, 51, 47, 58): 'n_total'}, {}), '(sentence_embedding, n_total)', True, 'import tensorflow as tf\n'), ((48, 15, 49, 57), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (), '', True, 'import tensorflow as tf\n'), ((50, 9, 50, 66), 'tensorflow.split', 'tf.split', ({(50, 18, 50, 28): 'all_params', (50, 30, 50, 62): '[(c * 2) for c in n_layer_channel]', (50, 64, 50, 65): '(1)'}, {}), '(all_params, [(c * 2) for c in n_layer_channel], 1)', True, 'import tensorflow as tf\n'), ((77, 9, 77, 36), 'tensorflow.keras.Sequential', 'tf.keras.Sequential', ({(77, 29, 77, 35): 'layers'}, {}), '(layers)', True, 'import tensorflow as tf\n'), ((93, 9, 93, 36), 'tensorflow.keras.Sequential', 'tf.keras.Sequential', ({(93, 29, 93, 35): 'layers'}, {}), '(layers)', True, 'import tensorflow as tf\n'), ((126, 16, 126, 48), 'tensorflow.keras.layers.Lambda', 'tf.keras.layers.Lambda', ({(126, 39, 126, 47): 'tf.shape'}, {}), '(tf.shape)', True, 'import tensorflow as tf\n'), ((130, 24, 130, 73), 'tensorflow.keras.layers.Reshape', 'tf.keras.layers.Reshape', ({(130, 48, 130, 72): '(-1, 1, shape_numpy[-1])'}, {}), '((-1, 1, shape_numpy[-1]))', True, 'import tensorflow as tf\n'), ((131, 24, 131, 73), 'tensorflow.keras.layers.Reshape', 'tf.keras.layers.Reshape', ({(131, 48, 131, 72): '(-1, 1, shape_numpy[-1])'}, {}), '((-1, 1, shape_numpy[-1]))', True, 'import tensorflow as tf\n'), ((139, 9, 139, 44), 'tensorflow.keras.layers.concatenate', 'tf.keras.layers.concatenate', ({(139, 37, 139, 43): '[a, b]'}, {}), '([a, b])', True, 'import tensorflow as tf\n'), ((75, 6, 76, 68), 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (), '', True, 'import tensorflow as tf\n'), ((92, 16, 92, 52), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', ({(92, 38, 92, 51): 'layer_cfg[-1]'}, {}), '(layer_cfg[-1])', True, 'import tensorflow as tf\n'), ((138, 6, 138, 40), 'tensorflow.keras.layers.Permute', 'tf.keras.layers.Permute', ({(138, 30, 138, 39): '(2, 1, 3)'}, {}), '((2, 1, 3))', True, 'import tensorflow as tf\n'), ((158, 14, 158, 42), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', ({(158, 36, 158, 41): 'units'}, {}), '(units)', True, 'import tensorflow as tf\n'), ((159, 14, 159, 42), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', ({(159, 36, 159, 41): 'units'}, {}), '(units)', True, 'import tensorflow as tf\n'), ((160, 13, 160, 37), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', ({(160, 35, 160, 36): '1'}, {}), '(1)', True, 'import tensorflow as tf\n'), ((167, 28, 167, 53), 'tensorflow.expand_dims', 'tf.expand_dims', ({(167, 43, 167, 49): 'hidden', (167, 51, 167, 52): '1'}, {}), '(hidden, 1)', True, 'import tensorflow as tf\n'), ((178, 21, 178, 58), 'tensorflow.reduce_sum', 'tf.reduce_sum', (), '', True, 'import tensorflow as tf\n'), ((202, 21, 202, 73), 'tensorflow.keras.layers.Embedding', 'tf.keras.layers.Embedding', ({(202, 47, 202, 57): 'vocab_size', (202, 59, 202, 72): 'embedding_dim'}, {}), '(vocab_size, embedding_dim)', True, 'import tensorflow as tf\n'), ((203, 15, 207, 47), 'tensorflow.keras.layers.GRU', 'tf.keras.layers.GRU', (), '', True, 'import tensorflow as tf\n'), ((217, 11, 217, 46), 'tensorflow.zeros', 'tf.zeros', ({(217, 20, 217, 45): '(batch_size, self._units)'}, {}), '((batch_size, self._units))', True, 'import tensorflow as tf\n'), ((67, 8, 72, 28), 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (), '', True, 'import tensorflow as tf\n'), ((91, 18, 91, 67), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (), '', True, 'import tensorflow as tf\n'), ((135, 21, 135, 60), 'tensorflow.tile', 'tf.tile', (), '', True, 'import tensorflow as tf\n')]