# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import tempfile
import unittest
from pathlib import Path
from typing import Any, Dict, Sequence

import fairseq.data.indexed_dataset as indexed_dataset
import fairseq.options
import fairseq.tasks.online_backtranslation as obt
import torch
from tests import utils


def mk_sample(tokens: Sequence[int], batch_size: int = 2) -> Dict[str, Any]:
    batch = torch.stack([torch.tensor(tokens, dtype=torch.long)] * batch_size)
    sample = {
        "net_input": {
            "src_tokens": batch,
            "prev_output_tokens": batch,
            "src_lengths": torch.tensor([len(tokens)] * batch_size, dtype=torch.long),
        },
        "target": batch[:, 1:],
    }
    return sample


def mk_dataset(num_samples: int, max_len: int, output: Path):
    output.parent.mkdir(exist_ok=True)
    idx = indexed_dataset.IndexedDatasetBuilder(str(output))
    data = torch.randint(5, 100, (num_samples, max_len))
    lengths = torch.randint(3, max_len, (num_samples,))
    for d, l in zip(data, lengths):
        d[0] = 0
        idx.add_item(d[:l])
    idx.finalize(output.with_suffix(".idx"))
    assert output.exists()
    assert output.with_suffix(".idx").exists()


class OnlineBacktranslationTest(unittest.TestCase):

    tmp_dir = Path(tempfile.mkdtemp(suffix="OnlineBacktranslationTest"))

    @classmethod
    def obt_task(
        cls, languages: Sequence[str], data: Path = None, language_mapping: str = None
    ):
        dict_path = cls.tmp_dir / "dict.txt"
        if not dict_path.exists():
            dictionary = utils.dummy_dictionary(100)
            dictionary.save(str(dict_path))

        if data is not None:
            (data / "dict.txt").write_text(dict_path.read_text())
        else:
            data = cls.tmp_dir
        assert len(languages) >= 2

        kwargs = {
            "arch": "transformer",
            # --max-sentences=1 for better predictability of batches
            "max_sentences": 1,
            # Use characteristics dimensions
            "encoder_layers": 3,
            "encoder_embed_dim": 12,
            "encoder_ffn_embed_dim": 14,
            "encoder_attention_heads": 4,
            "decoder_layers": 3,
            "decoder_embed_dim": 12,
            "decoder_output_dim": 12,
            "decoder_ffn_embed_dim": 14,
            "decoder_attention_heads": 4,
            # Disable dropout so we have comparable tests.
            "dropout": 0,
            "attention_dropout": 0,
            "activation_dropout": 0,
            "encoder_layerdrop": 0,
        }

        args = fairseq.options.get_args(
            data,
            task="online_backtranslation",
            mono_langs=",".join(languages),
            valid_lang_pairs=f"{languages[0]}-{languages[1]}",
            tokens_per_sample=256,
            language_mapping=language_mapping,
            **kwargs,
        )
        task = obt.OnlineBackTranslationTask.setup_task(args)
        # we need to build the model to have the correct dictionary
        model = task.build_model(task.args)
        return task, model

    def tmp_path(self, test_case: str) -> Path:
        return Path(tempfile.mkdtemp(test_case, dir=self.tmp_dir))

    def test_lang_tokens(self):
        task, model = self.obt_task(["en", "ro", "zh"])
        assert obt._lang_token("en") in task.dictionary
        assert obt._lang_token("ro") in task.dictionary
        assert obt._lang_token("zh") in task.dictionary

        en_bos = obt._lang_token_index(task.common_dict, "en")
        assert "en" == task.common_dict[en_bos].strip("_")
        zh_bos = obt._lang_token_index(task.common_dict, "zh")
        assert "zh" == task.common_dict[zh_bos].strip("_")
        zh_sample = mk_sample([zh_bos, 16, 14, 12, 10])

        # we expect to receive the bos token for translation
        assert task.get_bos_token_from_sample(zh_sample) == en_bos

    def test_backtranslate_sample(self):
        task, model = self.obt_task(["en", "ro", "zh"])

        en_bos = obt._lang_token_index(task.common_dict, "en")
        zh_bos = obt._lang_token_index(task.common_dict, "zh")
        sample = mk_sample([zh_bos, 16, 14, 12, 10])

        task.backtranslate_sample(sample, "zh", "en")
        target_zh = list(sample["target"][0])
        assert target_zh == [16, 14, 12, 10]  # original zh sentence
        generated_en = sample["net_input"]["src_tokens"][0]
        assert generated_en[0] == en_bos

    def test_train_dataset(self):
        data = self.tmp_path("test_train_dataset")
        mk_dataset(20, 10, data / "en" / "train.bin")
        mk_dataset(10, 10, data / "zh" / "train.bin")
        task, model = self.obt_task(["en", "zh"], data)
        task.load_dataset("train")

        en_bos = obt._lang_token_index(task.common_dict, "en")
        zh_bos = obt._lang_token_index(task.common_dict, "zh")

        train = task.datasets["train"]
        train.ordered_indices()
        train.prefetch([0, 19])
        sample_0 = train[0]
        sample_19 = train[19]
        self.assertEqual(
            set(sample_0.keys()), {"en-BT", "en-DENOISE", "zh-BT", "zh-DENOISE"}
        )
        for sample in (sample_0, sample_19):
            self.assertEqual(sample["en-BT"]["source"][0], en_bos)
            # bt target isn't ready to look at.
            self.assertEqual(sample["en-DENOISE"]["source"][0], en_bos)
            # TODO What could we check on the target side ?

        for i in range(10):
            # Zh dataset is shorter, and is wrapped around En dataset.
            train.prefetch([i, i + 10])
            self.assertEqual(
                list(train[i]["zh-DENOISE"]["source"]),
                list(train[i + 10]["zh-DENOISE"]["source"]),
            )
            self.assertEqual(train[i]["zh-DENOISE"]["source"][0].item(), zh_bos)

        # Sorted by increasing len
        self.assertLess(
            len(sample_0["en-BT"]["source"]), len(sample_19["en-BT"]["source"])
        )

    def test_valid_dataset(self):
        data = self.tmp_path("test_valid_dataset")
        mk_dataset(10, 21, data / "valid.en-zh.en.bin")
        mk_dataset(10, 21, data / "valid.en-zh.zh.bin")

        task, model = self.obt_task(["en", "zh"], data)
        valid = task.load_dataset("valid")
        en_bos = obt._lang_token_index(task.common_dict, "en")

        assert valid is not None
        valid.prefetch(range(10))
        sample_0 = valid[0]
        sample_9 = valid[9]
        self.assertEqual(sample_0["id"], 0)
        self.assertEqual(sample_9["id"], 9)
        self.assertEqual(sample_0["source"][0], en_bos)
        self.assertEqual(sample_9["source"][0], en_bos)
        # TODO: could we test the target side ?

    def assertFnMatch(self, fn, values):
        for x, y in values.items():
            fn_x = fn(x)
            self.assertEqual(fn_x, y, f"Fn has wrong value: fn({x}) = {fn_x} != {y}")

    def test_piecewise_linear_fn(self):
        self.assertFnMatch(
            obt.PiecewiseLinearFn.from_string("1.0"), {0: 1, 100: 1, 500: 1, 1000: 1}
        )
        self.assertFnMatch(
            obt.PiecewiseLinearFn.from_string("0:1,1000:0"),
            {0: 1, 500: 0.5, 1000: 0, 2000: 0},
        )
        self.assertFnMatch(
            obt.PiecewiseLinearFn.from_string("0:0,1000:1"),
            {0: 0, 500: 0.5, 1000: 1, 2000: 1},
        )
        self.assertFnMatch(
            obt.PiecewiseLinearFn.from_string("0:0,1000:1,2000:0"),
            {0: 0, 500: 0.5, 1000: 1, 1500: 0.5, 2000: 0, 3000: 0},
        )