# Copyright 2019 The Texar Authors. All Rights Reserved.
#
# 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.
"""
Unit tests for data related operations.
"""
import tempfile
import copy
import numpy as np

import unittest

from texar.torch.data.data.data_iterators import DataIterator
from texar.torch.data.data.paired_text_data import PairedTextData
from texar.torch.data.vocabulary import SpecialTokens


class PairedTextDataTest(unittest.TestCase):
    """Tests paired text data class.
    """

    def setUp(self):

        # Create test data
        vocab_list = ['This', 'is', 'a', 'word', '词']
        vocab_file = tempfile.NamedTemporaryFile()
        vocab_file.write('\n'.join(vocab_list).encode("utf-8"))
        vocab_file.flush()
        self._vocab_file = vocab_file
        self._vocab_size = len(vocab_list)

        src_text = ['This is a sentence from source .', '词 词 。 source']
        src_text_file = tempfile.NamedTemporaryFile()
        src_text_file.write('\n'.join(src_text).encode("utf-8"))
        src_text_file.flush()
        self._src_text_file = src_text_file

        tgt_text = ['This is a sentence from target .', '词 词 。 target']
        tgt_text_file = tempfile.NamedTemporaryFile()
        tgt_text_file.write('\n'.join(tgt_text).encode("utf-8"))
        tgt_text_file.flush()
        self._tgt_text_file = tgt_text_file

        self._hparams = {
            "num_epochs": 1,
            "batch_size": 3,
            "source_dataset": {
                "files": [self._src_text_file.name],
                "vocab_file": self._vocab_file.name,
            },
            "target_dataset": {
                "files": self._tgt_text_file.name,
                "vocab_share": True,
                "eos_token": "<TARGET_EOS>"
            }
        }

        self.src_upper_cased_text = []
        for sent in src_text:
            upper_cased_tokens = sent.upper().split(" ")
            upper_cased_tokens.append(SpecialTokens.EOS)
            self.src_upper_cased_text.append(upper_cased_tokens)

        max_length = max([len(tokens) for tokens in self.src_upper_cased_text])
        self.src_upper_cased_text = [sent + [''] * (max_length - len(sent)) for
                                     sent in self.src_upper_cased_text]

        self.tgt_upper_cased_text = []
        for sent in tgt_text:
            upper_cased_tokens = sent.upper().split(" ")
            upper_cased_tokens.insert(0, "<BOS>")
            upper_cased_tokens.append("<TARGET_EOS>")
            self.tgt_upper_cased_text.append(upper_cased_tokens)

        max_length = max([len(tokens) for tokens in self.tgt_upper_cased_text])
        self.tgt_upper_cased_text = [sent + [''] * (max_length - len(sent)) for
                                     sent in self.tgt_upper_cased_text]

    def _run_and_test(self, hparams, proc_shr=False, test_transform=None,
                      discard_src=False):
        # Construct database
        text_data = PairedTextData(hparams)
        self.assertEqual(text_data.source_vocab.size,
                         self._vocab_size +
                         len(text_data.source_vocab.special_tokens))

        iterator = DataIterator(text_data)
        for data_batch in iterator:
            self.assertEqual(set(data_batch.keys()),
                             set(text_data.list_items()))

            if proc_shr:
                tgt_eos = '<EOS>'
            else:
                tgt_eos = '<TARGET_EOS>'

            # Test matching
            src_text = data_batch['source_text']
            tgt_text = data_batch['target_text']
            if proc_shr:
                for src, tgt in zip(src_text, tgt_text):
                    np.testing.assert_array_equal(src[:3], tgt[:3])
            else:
                for src, tgt in zip(src_text, tgt_text):
                    np.testing.assert_array_equal(src[:3], tgt[1:4])
            self.assertTrue(
                tgt_eos in data_batch['target_text'][0])

            if test_transform:
                for i in range(len(data_batch['source_text'])):
                    text_ = data_batch['source_text'][i]
                    self.assertTrue(text_ in self.src_upper_cased_text)
                for i in range(len(data_batch['target_text'])):
                    text_ = data_batch['target_text'][i]
                    self.assertTrue(text_ in self.tgt_upper_cased_text)

            if discard_src:
                src_hparams = text_data.hparams.source_dataset
                max_l = src_hparams.max_seq_length
                max_l += sum(int(x is not None and x != '')
                             for x in [text_data._src_bos_token,
                                       text_data._tgt_bos_token])
                for l in data_batch["source_length"]:
                    self.assertLessEqual(l, max_l)

    def test_default_setting(self):
        """Tests the logics of the text data.
        """
        self._run_and_test(self._hparams)

    def test_shuffle(self):
        """Tests toggling shuffle.
        """
        hparams = copy.copy(self._hparams)
        hparams["shuffle"] = False
        self._run_and_test(hparams)

    def test_processing_share(self):
        """Tests sharing processing.
        """
        hparams = copy.copy(self._hparams)
        hparams["target_dataset"]["processing_share"] = True
        self._run_and_test(hparams, proc_shr=True)

    def test_other_transformations(self):
        """Tests use of other transformations
        """
        _upper_func = lambda raw_example: [x.upper() for x in raw_example]

        hparams = copy.copy(self._hparams)
        hparams["source_dataset"].update(
            {"other_transformations": [_upper_func]})
        hparams["target_dataset"].update(
            {"other_transformations": [_upper_func]})
        self._run_and_test(hparams, test_transform=True)

    def test_length_filter(self):
        """Tests filtering by length.
        """
        hparams = copy.copy(self._hparams)
        hparams["source_dataset"].update(
            {"max_seq_length": 4,
             "length_filter_mode": "discard"})
        self._run_and_test(hparams, discard_src=True)


if __name__ == "__main__":
    unittest.main()
