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


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl.testing import parameterized

import mesh_tensorflow as mtf
from mesh_tensorflow.transformer import utils

import tensorflow.compat.v1 as tf

tf.disable_v2_behavior()


class UtilsTest(parameterized.TestCase, tf.test.TestCase):

  def testDynamicText2self(self):
    batch = 2
    length = 5
    input_tensors = {
        "inputs": [
            [3, 1, 4, 1, 0],
            [1, 4, 3, 2, 1]],
        "inputs_segmentation": [
            [1, 1, 2, 2, 0],
            [1, 2, 2, 2, 2]],
        "inputs_position": [
            [0, 1, 0, 1, 0],
            [0, 0, 1, 2, 3]],
        "targets": [
            [1, 1, 0, 0, 0],
            [9, 8, 1, 2, 1]],
        "targets_segmentation": [
            [1, 2, 0, 0, 0],
            [1, 1, 1, 2, 2]],
        "targets_position": [
            [0, 0, 0, 0, 0],
            [0, 1, 2, 0, 1]]
    }
    expected_output_tensors = {
        "targets": [
            [3, 1, 1, 4, 1, 1, 0, 0, 0, 0],
            [1, 9, 8, 1, 4, 3, 2, 1, 2, 1]],
        "targets_segmentation": [
            [1, 1, 1, 2, 2, 2, 0, 0, 0, 0],
            [1, 1, 1, 1, 2, 2, 2, 2, 2, 2]],
        "targets_position": [
            [0, 1, 2, 0, 1, 2, 0, 0, 0, 0],
            [0, 1, 2, 3, 0, 1, 2, 3, 4, 5]]
    }
    graph = mtf.Graph()
    mesh = mtf.Mesh(graph, "my_mesh")
    batch_dim = mtf.Dimension("batch", batch)
    length_dim = mtf.Dimension("length", length)

    input_shape = mtf.Shape([batch_dim, length_dim])
    mtf_features = {k: mtf.import_tf_tensor(mesh, v, input_shape)
                    for k, v in input_tensors.items()}
    mtf_outputs = utils._dynamic_text2self(mtf_features)
    mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
        shape=[], layout={}, devices=[""])
    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
    for k, v in expected_output_tensors.items():
      out = lowering.export_to_tf_tensor(mtf_outputs[k])
      actual = self.evaluate(out)
      self.assertAllEqual(actual, v)

  def testCleanDecodes(self):
    cleaned_decodes = utils.clean_decodes([
        [2, 0, 2, 1, 3, 2, 0],
        [1, 2, 2, 2, 2, 2, 2],
        [2, 2, 1, 1, 1, 1, 1],
        [2, 2, 2, 2, 2, 2, 2]
    ])
    with self.test_session() as sess:
      self.assertAllEqual(
          sess.run(cleaned_decodes),
          [
              [2, 0, 2, 1, 0, 0, 0],
              [1, 0, 0, 0, 0, 0, 0],
              [2, 2, 1, 0, 0, 0, 0],
              [2, 2, 2, 2, 2, 2, 2]
          ])

if __name__ == "__main__":
  tf.test.main()
