# coding=utf-8
# Copyright 2021 The Trax 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.

"""End to end test for Reformer."""

import os

from absl.testing import absltest
import gin

from trax import test_utils
from trax.models.reformer import reformer  # pylint: disable=unused-import
from trax.supervised import trainer_lib
from trax.tf_numpy import numpy as tf_np  # pylint: disable=unused-import

pkg_dir, _ = os.path.split(__file__)
_TESTDATA = os.path.join(pkg_dir, 'testdata')
_CONFIG_DIR = os.path.join(pkg_dir, '../../supervised/configs/')


class ReformerE2ETest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    gin.clear_config()
    gin.add_config_file_search_path(_CONFIG_DIR)
    test_utils.ensure_flag('test_tmpdir')

  def test_reformer_wmt_ende(self):
    batch_size_per_device = 2
    steps = 1
    n_layers = 2
    d_ff = 32

    gin.parse_config_file('reformer_wmt_ende.gin')

    gin.bind_parameter('data_streams.data_dir', _TESTDATA)
    gin.bind_parameter('batcher.batch_size_per_device', batch_size_per_device)
    gin.bind_parameter('train.steps', steps)
    gin.bind_parameter('Reformer.n_encoder_layers', n_layers)
    gin.bind_parameter('Reformer.n_decoder_layers', n_layers)
    gin.bind_parameter('Reformer.d_ff', d_ff)

    output_dir = self.create_tempdir().full_path
    _ = trainer_lib.train(output_dir=output_dir)

  def test_reformer_copy(self):
    batch_size_per_device = 2
    steps = 1
    n_layers = 2
    d_ff = 32
    d_model = 32

    gin.parse_config_file('reformer_copy.gin')

    gin.bind_parameter('data_streams.data_dir', _TESTDATA)
    gin.bind_parameter('batcher.batch_size_per_device', batch_size_per_device)
    gin.bind_parameter('train.steps', steps)
    gin.bind_parameter('ReformerLM.n_layers', n_layers)
    gin.bind_parameter('ReformerLM.d_ff', d_ff)
    gin.bind_parameter('ReformerLM.d_model', d_model)

    output_dir = self.create_tempdir().full_path
    _ = trainer_lib.train(output_dir=output_dir)


if __name__ == '__main__':
  absltest.main()
