# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.

import torch

from fairseq import sequence_generator
from fairseq import utils

from . import FairseqDataset, language_pair_dataset


class BacktranslationDataset(FairseqDataset):
    def __init__(
        self,
        tgt_dataset,
        tgt_dict,
        backtranslation_model,
        max_len_a,
        max_len_b,
        remove_eos_at_src=False,
        generator_class=sequence_generator.SequenceGenerator,
        cuda=True,
        **kwargs
    ):
        """
        Sets up a backtranslation dataset which takes a tgt batch, generates
        a src using a tgt-src backtranslation_model, and returns the
        corresponding {generated src, input tgt} batch
        Args:
            tgt_dataset: dataset which will be used to build self.tgt_dataset --
                a LanguagePairDataset with tgt dataset as the source dataset and
                None as the target dataset. Should NOT have padding so that
                src_lengths are accurately calculated by language_pair_dataset
                collate function.
                We use language_pair_dataset here to encapsulate the tgt_dataset
                so we can re-use the LanguagePairDataset collater to format the
                batches in the structure that SequenceGenerator expects.
                Note: tgt_dataset samples should not have EOS at end if
                the tgt-src model expects an input without EOS. This dataset
                does not enforce this, you should enforce that in preprocessing.
            tgt_dict: tgt dictionary (typically a joint src/tgt BPE dictionary)
            backtranslation_model: tgt-src model to use in the SequenceGenerator
                to generate backtranslations from tgt batches
            max_len_a, max_len_b: args passed into generate() function of
                the backtranslation SequenceGenerator
            remove_eos_at_src: whether we should remove EOS from the source
                dialect text generated by the backtranslation model.
            generator_class: which SequenceGenerator class to use for
                backtranslation. Output of generate() should be the same format
                as fairseq's SequenceGenerator
            cuda: use GPU for generation
            kwargs: generation args to init the backtranslation
                SequenceGenerator
        """
        self.tgt_dataset = language_pair_dataset.LanguagePairDataset(
            src=tgt_dataset,
            src_sizes=tgt_dataset.sizes,
            src_dict=tgt_dict,
            tgt=None,
            tgt_sizes=None,
            tgt_dict=None,
        )

        self.max_len_a = max_len_a
        self.max_len_b = max_len_b
        self.remove_eos_at_src = remove_eos_at_src

        self.backtranslation_generator = generator_class(
            models=[backtranslation_model],
            tgt_dict=tgt_dict,
            **kwargs,
        )

        self.cuda = cuda if torch.cuda.is_available() else False
        if self.cuda:
            self.backtranslation_generator.cuda()

    def __getitem__(self, index):
        """
        Returns a single sample. Multiple samples are fed to the collater to
        create a backtranslation batch. Note you should always use collate_fn
        BacktranslationDataset.collater() below if given the option to
        specify which collate_fn to use (e.g. in a dataloader which uses this
        BacktranslationDataset -- see corresponding unittest for an example).
        """
        return self.tgt_dataset[index]

    def __len__(self):
        """
        The length of the backtranslation dataset is the length of tgt.
        """
        return len(self.tgt_dataset)

    def collater(self, samples):
        """
        Using the samples from the tgt dataset, load a collated tgt sample to
        feed to the backtranslation model. Then take the generated translation
        with best score as the source and the orignal net input as the target.
        """
        collated_tgt_only_sample = self.tgt_dataset.collater(samples=samples)
        backtranslation_hypos = self._generate_hypotheses(
            sample=collated_tgt_only_sample
        )

        # Go through each tgt sentence in batch and its corresponding best
        # generated hypothesis and create a backtranslation data pair
        # {id: id, source: generated backtranslation, target: original tgt}
        generated_samples = []
        for input_sample, hypos in zip(samples, backtranslation_hypos):
            original_tgt = input_sample["source"].cpu()
            generated_source = hypos[0]["tokens"].cpu()  # first hypo is best hypo

            # Append EOS to the tgt sentence if it does not have an EOS
            # This is the case if the samples in monolingual tgt_dataset don't
            # have an EOS appended to the end of each sentence.
            eos = self.tgt_dataset.src_dict.eos()
            if original_tgt[-1] != eos:
                original_tgt = torch.cat([original_tgt, torch.LongTensor([eos])])

            # The generated source dialect backtranslation will have an EOS.
            # If we want our parallel data source to not have an EOS, we will
            # have to remove it.
            if self.remove_eos_at_src:
                assert generated_source[-1] == eos, (
                    "Expected generated backtranslation to have eos (id: "
                    "{eos}) at end, but instead found token id "
                    "{generated_source[-1]} at end."
                ).format(eos=eos, generated_source=generated_source)
                generated_source = generated_source[:-1]

            generated_samples.append(
                {
                    "id": input_sample["id"],
                    "source": generated_source,
                    "target": original_tgt,
                }
            )

        return language_pair_dataset.collate(
            samples=generated_samples,
            pad_idx=self.tgt_dataset.src_dict.pad(),
            eos_idx=self.tgt_dataset.src_dict.eos(),
        )

    def get_dummy_batch(self, num_tokens, max_positions):
        """ Just use the tgt dataset get_dummy_batch """
        return self.tgt_dataset.get_dummy_batch(num_tokens, max_positions)

    def num_tokens(self, index):
        """ Just use the tgt dataset num_tokens """
        return self.tgt_dataset.num_tokens(index)

    def ordered_indices(self):
        """ Just use the tgt dataset ordered_indices """
        return self.tgt_dataset.ordered_indices()

    def valid_size(self, index, max_positions):
        """ Just use the tgt dataset size """
        return self.tgt_dataset.valid_size(index, max_positions)

    def _generate_hypotheses(self, sample):
        """
        Generates hypotheses from a LanguagePairDataset collated / batched
        sample. Note in this case, sample["target"] is None, and
        sample["net_input"]["src_tokens"] is really in tgt language.
        """
        s = utils.move_to_cuda(sample) if self.cuda else sample
        input = s["net_input"]
        srclen = input["src_tokens"].size(1)
        hypos = self.backtranslation_generator.generate(
            input,
            maxlen=int(
                self.max_len_a * srclen + self.max_len_b
            ),
        )
        return hypos

    def size(self, index):
        """Return an example's size as a float or tuple. This value is used when
        filtering a dataset with ``--max-positions``.
        Here, we return src dataset size as tgt dataset size as an approximation.
        We do not know src size until we backtranslate and generate src sentences.
        """
        return (self.tgt_dataset.size(index)[0], self.tgt_dataset.size(index)[0])
