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

"""Tests for Transformer."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np

from tensor2tensor.data_generators import librispeech
from tensor2tensor.data_generators import problem_hparams
from tensor2tensor.models import transformer

import tensorflow.compat.v1 as tf


BATCH_SIZE = 3
INPUT_LENGTH = 5
TARGET_LENGTH = 7
VOCAB_SIZE = 10


def get_model(hparams=None, mode=tf.estimator.ModeKeys.TRAIN,
              has_input=True, model_cls=transformer.Transformer):
  if hparams is None:
    hparams = transformer.transformer_tiny()
  hparams.hidden_size = 8
  hparams.filter_size = 32
  hparams.num_heads = 1
  hparams.layer_prepostprocess_dropout = 0.0

  if hparams.get("problem_hparams", None) is None:
    p_hparams = problem_hparams.test_problem_hparams(VOCAB_SIZE,
                                                     VOCAB_SIZE,
                                                     hparams)
  if not has_input:
    del p_hparams.modality["inputs"]
  hparams.problem_hparams = p_hparams

  inputs = np.random.randint(
      VOCAB_SIZE, size=(BATCH_SIZE, INPUT_LENGTH, 1, 1))
  targets = np.random.randint(
      VOCAB_SIZE, size=(BATCH_SIZE, TARGET_LENGTH, 1, 1))
  features = {
      "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
      "target_space_id": tf.constant(1, dtype=tf.int32)
  }
  if has_input:
    features["inputs"] = tf.constant(inputs, dtype=tf.int32, name="inputs")

  return model_cls(hparams, mode, p_hparams), features


def small_librispeech_model(param_overrides=None):
  hparams = transformer.transformer_small()
  hparams.hidden_size = 8
  hparams.filter_size = 32
  hparams.num_heads = 1
  hparams.layer_prepostprocess_dropout = 0.0
  p_hparams = librispeech.Librispeech().get_hparams(hparams)
  p_hparams.vocab_size["targets"] = VOCAB_SIZE
  hparams.problem_hparams = p_hparams
  model = transformer.Transformer(hparams, problem_hparams=p_hparams)
  if param_overrides is not None:  # Add or Set any provided HParams
    assert isinstance(param_overrides, dict)
    for param_name in param_overrides:
      if hasattr(hparams, param_name):
        hparams.set_hparam(param_name, param_overrides[param_name])
      else:
        hparams.add_hparam(param_name, param_overrides[param_name])
  inputs = np.random.rand(
      BATCH_SIZE, INPUT_LENGTH, 80, 3).astype("float32")  # modify for speech
  targets = np.random.randint(
      VOCAB_SIZE, size=(BATCH_SIZE, TARGET_LENGTH, 1, 1))
  features = {
      "inputs": tf.constant(inputs, dtype=tf.float32, name="inputs"),
      "targets": tf.constant(targets, dtype=tf.int32, name="targets"),
      "target_space_id": tf.constant(1, dtype=tf.int32)
  }
  return model, features


class TransformerTest(tf.test.TestCase):

  def testTransformer(self, get_model_fn=None, p=None):
    if get_model_fn:
      model, features = get_model_fn(param_overrides=p)
    else:
      model, features = get_model(transformer.transformer_small())
    logits, _ = model(features)
    with self.test_session() as session:
      session.run(tf.global_variables_initializer())
      res = session.run(logits)
    self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, 1, 1, VOCAB_SIZE))

  def testTransformerLibrispeech(self, params=None):
    self.testTransformer(get_model_fn=small_librispeech_model, p=params)

  def testLibrispeechSlowVsFast(self, params=None):
    self.testSlowVsFast(get_model_fn=small_librispeech_model, p=params)

  def testLibrispeechMultihead(self, params=None):
    self.testTransformerLibrispeech({"num_heads": 2})

  def testLibrispeechWithAreaAttention(self):
    self.testTransformerLibrispeech({"max_area_width": 2,
                                     "num_area_layers": 1,
                                     "area_key_mode": "mean",
                                     "area_value_mode": "sum"})

  def testTransformerRelative(self):
    model, features = get_model(transformer.transformer_relative_tiny())
    logits, _ = model(features)
    with self.test_session() as session:
      session.run(tf.global_variables_initializer())
      res = session.run(logits)
    self.assertEqual(res.shape, (BATCH_SIZE, TARGET_LENGTH, 1, 1, VOCAB_SIZE))

  def testSlowVsFast(self, get_model_fn=None, p=None):
    if get_model_fn:
      model, features = get_model_fn(param_overrides=p)
    else:
      model, features = get_model(transformer.transformer_small())

    decode_length = 3

    out_logits, _ = model(features)
    out_logits = tf.squeeze(out_logits, axis=[2, 3])
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(