# 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 torch

from . import FairseqDataset


class TransformEosDataset(FairseqDataset):
    """A :class:`~fairseq.data.FairseqDataset` wrapper that appends/prepends/strips EOS.

    Note that the transformation is applied in :func:`collater`.

    Args:
        dataset (~fairseq.data.FairseqDataset): dataset to wrap
        eos (int): index of the end-of-sentence symbol
        append_eos_to_src (bool, optional): append EOS to the end of src
        remove_eos_from_src (bool, optional): remove EOS from the end of src
        append_eos_to_tgt (bool, optional): append EOS to the end of tgt
        remove_eos_from_tgt (bool, optional): remove EOS from the end of tgt
    """

    def __init__(
        self,
        dataset,
        eos,
        append_eos_to_src=False,
        remove_eos_from_src=False,
        append_eos_to_tgt=False,
        remove_eos_from_tgt=False,
        has_target=True,
    ):
        if not isinstance(dataset, FairseqDataset):
            raise ValueError('dataset must be an instance of FairseqDataset')
        if append_eos_to_src and remove_eos_from_src:
            raise ValueError('cannot combine append_eos_to_src and remove_eos_from_src')
        if append_eos_to_tgt and remove_eos_from_tgt:
            raise ValueError('cannot combine append_eos_to_tgt and remove_eos_from_tgt')

        self.dataset = dataset
        self.eos = torch.LongTensor([eos])
        self.append_eos_to_src = append_eos_to_src
        self.remove_eos_from_src = remove_eos_from_src
        self.append_eos_to_tgt = append_eos_to_tgt
        self.remove_eos_from_tgt = remove_eos_from_tgt
        self.has_target = has_target

        # precompute how we should adjust the reported sizes
        self._src_delta = 0
        self._src_delta += 1 if append_eos_to_src else 0
        self._src_delta -= 1 if remove_eos_from_src else 0
        self._tgt_delta = 0
        self._tgt_delta += 1 if append_eos_to_tgt else 0
        self._tgt_delta -= 1 if remove_eos_from_tgt else 0

        self._checked_src = False
        self._checked_tgt = False

    def _check_src(self, src, expect_eos):
        if not self._checked_src:
            assert (src[-1] == self.eos[0]) == expect_eos
            self._checked_src = True

    def _check_tgt(self, tgt, expect_eos):
        if self.has_target and not self._checked_tgt:
            assert (tgt[-1] == self.eos[0]) == expect_eos
            self._checked_tgt = True

    def __getitem__(self, index):
        return self.dataset[index]

    def __len__(self):
        return len(self.dataset)

    def collater(self, samples):

        def transform(item):
            if self.append_eos_to_src:
                self.eos = self.eos.to(device=item['source'].device)
                self._check_src(item['source'], expect_eos=False)
                item['source'] = torch.cat([item['source'], self.eos])
            if self.remove_eos_from_src:
                self.eos = self.eos.to(device=item['source'].device)
                self._check_src(item['source'], expect_eos=True)
                item['source'] = item['source'][:-1]
            if self.append_eos_to_tgt:
                self.eos = self.eos.to(device=item['target'].device)
                self._check_tgt(item['target'], expect_eos=False)
                item['target'] = torch.cat([item['target'], self.eos])
            if self.remove_eos_from_tgt:
                self.eos = self.eos.to(device=item['target'].device)
                self._check_tgt(item['target'], expect_eos=True)
                item['target'] = item['target'][:-1]
            return item

        samples = list(map(transform, samples))
        return self.dataset.collater(samples)

    def num_tokens(self, index):
        return self.dataset.num_tokens(index)

    def size(self, index):
        if self.has_target:
            src_len, tgt_len = self.dataset.size(index)
            return (src_len + self._src_delta, tgt_len + self._tgt_delta)
        else:
            return self.dataset.size(index)

    def ordered_indices(self):
        # NOTE: we assume that the ordering does not change based on the
        # addition or removal of eos
        return self.dataset.ordered_indices()

    @property
    def supports_prefetch(self):
        return getattr(self.dataset, 'supports_prefetch', False)

    def prefetch(self, indices):
        return self.dataset.prefetch(indices)
