# Copyright 2020 The HuggingFace Team. 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.

import os
from pathlib import Path

import numpy as np
import pytest
from pack_dataset import pack_data_dir
from parameterized import parameterized
from save_len_file import save_len_file
from torch.utils.data import DataLoader

from transformers import AutoTokenizer
from transformers.models.mbart.modeling_mbart import shift_tokens_right
from transformers.testing_utils import TestCasePlus, slow
from utils import FAIRSEQ_AVAILABLE, DistributedSortishSampler, LegacySeq2SeqDataset, Seq2SeqDataset


BERT_BASE_CASED = "bert-base-cased"
PEGASUS_XSUM = "google/pegasus-xsum"
ARTICLES = [" Sam ate lunch today.", "Sams lunch ingredients."]
SUMMARIES = ["A very interesting story about what I ate for lunch.", "Avocado, celery, turkey, coffee"]
T5_TINY = "patrickvonplaten/t5-tiny-random"
BART_TINY = "sshleifer/bart-tiny-random"
MBART_TINY = "sshleifer/tiny-mbart"
MARIAN_TINY = "sshleifer/tiny-marian-en-de"


def _dump_articles(path: Path, articles: list):
    content = "\n".join(articles)
    Path(path).open("w").writelines(content)


def make_test_data_dir(tmp_dir):
    for split in ["train", "val", "test"]:
        _dump_articles(os.path.join(tmp_dir, f"{split}.source"), ARTICLES)
        _dump_articles(os.path.join(tmp_dir, f"{split}.target"), SUMMARIES)
    return tmp_dir


class TestAll(TestCasePlus):
    @parameterized.expand(
        [
            MBART_TINY,
            MARIAN_TINY,
            T5_TINY,
            BART_TINY,
            PEGASUS_XSUM,
        ],
    )
    @slow
    def test_seq2seq_dataset_truncation(self, tok_name):
        tokenizer = AutoTokenizer.from_pretrained(tok_name)
        tmp_dir = make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())
        max_len_source = max(len(tokenizer.encode(a)) for a in ARTICLES)
        max_len_target = max(len(tokenizer.encode(a)) for a in SUMMARIES)
        max_src_len = 4
        max_tgt_len = 8
        assert max_len_target > max_src_len  # Will be truncated
        assert max_len_source > max_src_len  # Will be truncated
        src_lang, tgt_lang = "ro_RO", "de_DE"  # ignored for all but mbart, but never causes error.
        train_dataset = Seq2SeqDataset(
            tokenizer,
            data_dir=tmp_dir,
            type_path="train",
            max_source_length=max_src_len,
            max_target_length=max_tgt_len,  # ignored
            src_lang=src_lang,
            tgt_lang=tgt_lang,
        )
        dataloader = DataLoader(train_dataset, batch_size=2, collate_fn=train_dataset.collate_fn)
        for batch in dataloader:
            assert isinstance(batch, dict)
            assert batch["attention_mask"].shape == batch["input_ids"].shape
            # show that articles were trimmed.
            assert batch["input_ids"].shape[1] == max_src_len
            # show that targets are the same len
            assert batch["labels"].shape[1] == max_tgt_len
            if tok_name != MBART_TINY:
                continue
            # check language codes in correct place
            batch["decoder_input_ids"] = shift_tokens_right(batch["labels"], tokenizer.pad_token_id)
            assert batch["decoder_input_ids"][0, 0].item() == tokenizer.lang_code_to_id[tgt_lang]
            assert batch["decoder_input_ids"][0, -1].item() == tokenizer.eos_token_id
            assert batch["input_ids"][0, -2].item() == tokenizer.eos_token_id
            assert batch["input_ids"][0, -1].item() == tokenizer.lang_code_to_id[src_lang]

            break  # No need to test every batch

    @parameterized.expand([BART_TINY, BERT_BASE_CASED])
    def test_legacy_dataset_truncation(self, tok):
        tokenizer = AutoTokenizer.from_pretrained(tok)
        tmp_dir = make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())
        max_len_source = max(len(tokenizer.encode(a)) for a in ARTICLES)
        max_len_target = max(len(tokenizer.encode(a)) for a in SUMMARIES)
        trunc_target = 4
        train_dataset = LegacySeq2SeqDataset(
            tokenizer,
            data_dir=tmp_dir,
            type_path="train",
            max_source_length=20,
            max_target_length=trunc_target,
        )
        dataloader = DataLoader(train_dataset, batch_size=2, collate_fn=train_dataset.collate_fn)
        for batch in dataloader:
            assert batch["attention_mask"].shape == batch["input_ids"].shape
            # show that articles were trimmed.
            assert batch["input_ids"].shape[1] == max_len_source
            assert 20 >= batch["input_ids"].shape[1]  # trimmed significantly
            # show that targets were truncated
            assert batch["labels"].shape[1] == trunc_target  # Truncated
            assert max_len_target > trunc_target  # Truncated
            break  # No need to test every batch

    def test_pack_dataset(self):
        tokenizer = AutoTokenizer.from_pretrained("facebook/mbart-large-cc25")

        tmp_dir = Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()))
        orig_examples = tmp_dir.joinpath("train.source").open().readlines()
        save_dir = Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()))
        pack_data_dir(tokenizer, tmp_dir, 128, save_dir)
        orig_paths = {x.name for x in tmp_dir.iterdir()}
        new_paths = {x.name for x in save_dir.iterdir()}
        packed_examples = save_dir.joinpath("train.source").open().readlines()
        # orig: [' Sam ate lunch today.\n', 'Sams lunch ingredients.']
        # desired_packed: [' Sam ate lunch today.\n Sams lunch ingredients.']
        assert len(packed_examples) < len(orig_examples)
        assert len(packed_examples) == 1
        assert len(packed_examples[0]) == sum(len(x) for x in orig_examples)
        assert orig_paths == new_paths

    @pytest.mark.skipif(not FAIRSEQ_AVAILABLE, reason="This test requires fairseq")
    def test_dynamic_batch_size(self):
        if not FAIRSEQ_AVAILABLE:
            return
        ds, max_tokens, tokenizer = self._get_dataset(max_len=64)
        required_batch_size_multiple = 64
        batch_sampler = ds.make_dynamic_sampler(max_tokens, required_batch_size_multiple=required_batch_size_multiple)
        batch_sizes = [len(x) for x in batch_sampler]
        assert len(set(batch_sizes)) > 1  # it's not dynamic batch size if every batch is the same length
        assert sum(batch_sizes) == len(ds)  # no dropped or added examples
        data_loader = DataLoader(ds, batch_sampler=batch_sampler, collate_fn=ds.collate_fn, num_workers=2)
        failures = []
        num_src_per_batch = []
        for batch in data_loader:
            src_shape = batch["input_ids"].shape
            bs = src_shape[0]
            assert bs % required_batch_size_multiple == 0 or bs < required_batch_size_multiple
            num_src_tokens = np.product(batch["input_ids"].shape)
            num_src_per_batch.append(num_src_tokens)
            if num_src_tokens > (max_tokens * 1.1):
                failures.append(num_src_tokens)
        assert num_src_per_batch[0] == max(num_src_per_batch)
        if failures:
            raise AssertionError(f"too many tokens in {len(failures)} batches")

    def test_sortish_sampler_reduces_padding(self):
        ds, _, tokenizer = self._get_dataset(max_len=512)
        bs = 2
        sortish_sampler = ds.make_sortish_sampler(bs, shuffle=False)

        naive_dl = DataLoader(ds, batch_size=bs, collate_fn=ds.collate_fn, num_workers=2)
        sortish_dl = DataLoader(ds, batch_size=bs, collate_fn=ds.collate_fn, num_workers=2, sampler=sortish_sampler)

        pad = tokenizer.pad_token_id

        def count_pad_tokens(data_loader, k="input_ids"):
            return [batch[k].eq(pad).sum().item() for batch in data_loader]

        assert sum(count_pad_tokens(sortish_dl, k="labels")) < sum(count_pad_tokens(naive_dl, k="labels"))
        assert sum(count_pad_tokens(sortish_dl)) < sum(count_pad_tokens(naive_dl))
        assert len(sortish_dl) == len(naive_dl)

    def _get_dataset(self, n_obs=1000, max_len=128):
        if os.getenv("USE_REAL_DATA", False):
            data_dir = "examples/seq2seq/wmt_en_ro"
            max_tokens = max_len * 2 * 64
            if not Path(data_dir).joinpath("train.len").exists():
                save_len_file(MARIAN_TINY, data_dir)
        else:
            data_dir = "examples/seq2seq/test_data/wmt_en_ro"
            max_tokens = max_len * 4
            save_len_file(MARIAN_TINY, data_dir)

        tokenizer = AutoTokenizer.from_pretrained(MARIAN_TINY)
        ds = Seq2SeqDataset(
            tokenizer,
            data_dir=data_dir,
            type_path="train",
            max_source_length=max_len,
            max_target_length=max_len,
            n_obs=n_obs,
        )
        return ds, max_tokens, tokenizer

    def test_distributed_sortish_sampler_splits_indices_between_procs(self):
        ds, max_tokens, tokenizer = self._get_dataset()
        ids1 = set(DistributedSortishSampler(ds, 256, num_replicas=2, rank=0, add_extra_examples=False))
        ids2 = set(DistributedSortishSampler(ds, 256, num_replicas=2, rank=1, add_extra_examples=False))
        assert ids1.intersection(ids2) == set()

    @parameterized.expand(
        [
            MBART_TINY,
            MARIAN_TINY,
            T5_TINY,
            BART_TINY,
            PEGASUS_XSUM,
        ],
    )
    def test_dataset_kwargs(self, tok_name):
        tokenizer = AutoTokenizer.from_pretrained(tok_name, use_fast=False)
        if tok_name == MBART_TINY:
            train_dataset = Seq2SeqDataset(
                tokenizer,
                data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()),
                type_path="train",
                max_source_length=4,
                max_target_length=8,
                src_lang="EN",
                tgt_lang="FR",
            )
            kwargs = train_dataset.dataset_kwargs
            assert "src_lang" in kwargs and "tgt_lang" in kwargs
        else:
            train_dataset = Seq2SeqDataset(
                tokenizer,
                data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()),
                type_path="train",
                max_source_length=4,
                max_target_length=8,
            )
            kwargs = train_dataset.dataset_kwargs
            assert "add_prefix_space" not in kwargs if tok_name != BART_TINY else "add_prefix_space" in kwargs
            assert len(kwargs) == 1 if tok_name == BART_TINY else len(kwargs) == 0