# 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 copy
from argparse import Namespace
import json
import itertools
import logging
import os

import torch

from fairseq import options, search
import numpy as np
import torch.nn.functional as F

from fairseq import metrics, utils
from fairseq.criterions.my_at_criterion import label_smoothed_nll_loss
from fairseq.data import (
    AppendTokenDataset,
    ConcatDataset,
    data_utils,
    encoders,
    indexed_dataset,
    LanguagePairDataset,
    PrependTokenDataset,
    StripTokenDataset,
    TruncateDataset,
)

from fairseq import global_var_manager
from fairseq.tasks import register_task, LegacyFairseqTask
from fairseq.tasks.translation import load_langpair_dataset
from fairseq.utils import new_arange

EVAL_BLEU_ORDER = 4


logger = logging.getLogger(__name__)


def fetch_nat_input_indicator(at_probs, sample, strategy, threshold=None, lower_bound=0.0, upper_bound=1.0):
    if strategy == "at_below_threshold":
        assert 0 <= threshold <= 1
        at_golden_probs = torch.gather(at_probs, -1, sample['target'].unsqueeze(-1)).squeeze(-1)
        input_indicator = at_golden_probs > threshold
    elif strategy == "at_above_threshold":
        assert 0 <= threshold <= 1
        at_golden_probs = torch.gather(at_probs, -1, sample['target'].unsqueeze(-1)).squeeze(-1)
        input_indicator = at_golden_probs < threshold
    elif strategy == "confidence_interval":
        assert 0.0 <= lower_bound < upper_bound <= 1.0
        at_golden_probs = torch.gather(at_probs, -1, sample['target'].unsqueeze(-1)).squeeze(-1)
        input_indicator = (at_golden_probs <= lower_bound) | (at_golden_probs >= upper_bound)
    elif strategy == "at_wrong":
        at_pred_word_max_probs, at_pred_word_max_probs_indices = torch.max(at_probs, dim=-1)
        input_indicator = torch.eq(at_pred_word_max_probs_indices, sample['target'])
    elif strategy == "random":
        assert 0 <= threshold <= 1
        at_golden_probs = torch.gather(at_probs, -1, sample['target'].unsqueeze(-1)).squeeze(-1)
        input_indicator = at_golden_probs > threshold
        bsz, lenx = input_indicator.size()
        rand_index = torch.randperm(lenx, device=input_indicator.device).unsqueeze(0).expand(bsz, -1)
        input_indicator = torch.gather(input_indicator, -1, rand_index)
    elif strategy == "all-once":
        bsz, len_x, vocab = at_probs.size()
        temp_tensor = at_probs.new_zeros(bsz, len_x)
        input_indicator = temp_tensor != 0
    else:
        raise NotImplementedError

    return input_indicator


def fetch_at_nat_label_pred_sim(model, sample, at_probs, nat_probs, strategy=None):
    if strategy is None:
        return None
    elif strategy == "higher_golden_probs":
        at_golden_probs = torch.gather(at_probs, -1, sample['target'].unsqueeze(-1)).squeeze(-1)
        nat_golden_probs = torch.gather(nat_probs, -1, sample['target'].unsqueeze(-1)).squeeze(-1)
        label_pred_sim = (at_golden_probs, nat_golden_probs)
    elif strategy == "higher_emb_sim":
        bsz, len_x, vocab_size = at_probs.size()
        tgt_tokens = sample['target']
        # embed_matrix -> (B, V, d)
        temp_tensor = torch.arange(0, vocab_size, dtype=torch.long)
        at_embed_matrix = model.at_transformer.decoder.embed_tokens(
            temp_tensor.to(at_probs.device)).unsqueeze(0).expand(bsz, -1, -1)
        nat_embed_matrix = model.nat_transformer.decoder.embed_tokens(
            temp_tensor.to(at_probs.device)).unsqueeze(0).expand(bsz, -1, -1)
        # at_output_soft_embeddings -> (B, L, d)
        at_output_soft_embeddings = torch.bmm(at_probs, at_embed_matrix)
        nat_output_soft_embeddings = torch.bmm(nat_probs, nat_embed_matrix)
        # label_embeddings -> (B, L, d)
        at_label_embeddings = model.at_transformer.decoder.embed_tokens(tgt_tokens)
        nat_label_embeddings = model.nat_transformer.decoder.embed_tokens(tgt_tokens)
        # label_pred_sim: FloatTensor -> (B, L);
        at_label_pred_sim = F.cosine_similarity(at_output_soft_embeddings, at_label_embeddings,
                                                dim=-1)
        nat_label_pred_sim = F.cosine_similarity(nat_output_soft_embeddings, nat_label_embeddings,
                                                 dim=-1)
        # print("AT output_emb_sim: ", at_label_pred_sim)
        # print("NAT output_emb_sim: ", nat_label_pred_sim)
        label_pred_sim = (at_label_pred_sim, nat_label_pred_sim)
    elif strategy == "diff_on_nat_argmax":
        nat_argmax_probs, nat_argmax_prob_indices = nat_probs.max(dim=-1)
        at_prob_at_nat_argmax = torch.gather(at_probs, -1, nat_argmax_prob_indices.unsqueeze(-1)).squeeze(-1)
        label_pred_sim = (nat_argmax_probs, at_prob_at_nat_argmax)
    else:
        return NotImplementedError

    return label_pred_sim


@register_task('my_translation')
class MyTranslationTask(LegacyFairseqTask):

    @staticmethod
    def add_args(parser):
        """Add task-specific arguments to the parser."""
        # fmt: off
        parser.add_argument('data', help='colon separated path to data directories list, \
                                    will be iterated upon during epochs in round-robin manner; \
                                    however, valid and test data are always in the first directory to \
                                    avoid the need for repeating them in all directories')
        parser.add_argument('-s', '--source-lang', default=None, metavar='SRC',
                            help='source language')
        parser.add_argument('-t', '--target-lang', default=None, metavar='TARGET',
                            help='target language')
        parser.add_argument('--load-alignments', action='store_true',
                            help='load the binarized alignments')
        parser.add_argument('--left-pad-source', default='True', type=str, metavar='BOOL',
                            help='pad the source on the left')
        parser.add_argument('--left-pad-target', default='False', type=str, metavar='BOOL',
                            help='pad the target on the left')
        parser.add_argument('--max-source-positions', default=1024, type=int, metavar='N',
                            help='max number of tokens in the source sequence')
        parser.add_argument('--max-target-positions', default=1024, type=int, metavar='N',
                            help='max number of tokens in the target sequence')
        parser.add_argument('--upsample-primary', default=1, type=int,
                            help='amount to upsample primary dataset')
        parser.add_argument('--truncate-source', action='store_true', default=False,
                            help='truncate source to max-source-positions')
        parser.add_argument('--num-batch-buckets', default=0, type=int, metavar='N',
                            help='if >0, then bucket source and target lengths into N '
                                 'buckets and pad accordingly; this is useful on TPUs '
                                 'to minimize the number of compilations')

        # options for reporting BLEU during validation
        parser.add_argument('--eval-bleu', action='store_true',
                            help='evaluation with BLEU scores')
        parser.add_argument('--eval-bleu-detok', type=str, default="space",
                            help='detokenize before computing BLEU (e.g., "moses"); '
                                 'required if using --eval-bleu; use "space" to '
                                 'disable detokenization; see fairseq.data.encoders '
                                 'for other options')
        parser.add_argument('--eval-bleu-detok-args', type=str, metavar='JSON',
                            help='args for building the tokenizer, if needed')
        parser.add_argument('--eval-tokenized-bleu', action='store_true', default=False,
                            help='compute tokenized BLEU instead of sacrebleu')
        parser.add_argument('--eval-bleu-remove-bpe', nargs='?', const='@@ ', default=None,
                            help='remove BPE before computing BLEU')
        parser.add_argument('--eval-bleu-args', type=str, metavar='JSON',
                            help='generation args for BLUE scoring, '
                                 'e.g., \'{"beam": 4, "lenpen": 0.6}\'')
        parser.add_argument('--eval-bleu-print-samples', action='store_true',
                            help='print sample generations during validation')
        # options of NAT task (TranslationLevenshteinTask)
        parser.add_argument(
            '--noise',
            default='random_delete',
            choices=['random_delete', 'random_mask', 'no_noise', 'full_mask'])
        # fmt: on

    def __init__(self, args, src_dict, tgt_dict, mode):
        super().__init__(args)
        self.src_dict = src_dict
        self.tgt_dict = tgt_dict
        self.mode = mode

    @classmethod
    def setup_task(cls, args, **kwargs):
        """Setup the task (e.g., load dictionaries).

        Args:
            args (argparse.Namespace): parsed command-line arguments
        """
        args.left_pad_source = utils.eval_bool(args.left_pad_source)
        args.left_pad_target = utils.eval_bool(args.left_pad_target)

        paths = utils.split_paths(args.data)
        assert len(paths) > 0
        # find language pair automatically
        if args.source_lang is None or args.target_lang is None:
            args.source_lang, args.target_lang = data_utils.infer_language_pair(paths[0])
        if args.source_lang is None or args.target_lang is None:
            raise Exception('Could not infer language pair, please provide it explicitly')

        # load dictionaries
        src_dict = cls.load_dictionary(os.path.join(paths[0], 'dict.{}.txt'.format(args.source_lang)))
        tgt_dict = cls.load_dictionary(os.path.join(paths[0], 'dict.{}.txt'.format(args.target_lang)))
        assert src_dict.pad() == tgt_dict.pad()
        assert src_dict.eos() == tgt_dict.eos()
        assert src_dict.unk() == tgt_dict.unk()
        logger.info('[{}] dictionary: {} types'.format(args.source_lang, len(src_dict)))
        logger.info('[{}] dictionary: {} types'.format(args.target_lang, len(tgt_dict)))

        default_task_mode = getattr(args, "default_task_mode", None)
        task_mode = default_task_mode if default_task_mode is not None else getattr(args, "mode", "at")
        return cls(args, src_dict, tgt_dict, mode=task_mode)

    def load_dataset(self, split, epoch=1, combine=False, **kwargs):
        """Load a given dataset split.

        Args:
            split (str): name of the split (e.g., train, valid, test)
        """
        paths = utils.split_paths(self.args.data)
        assert len(paths) > 0
        if split != getattr(self.args, "train_subset", None):
            # if not training data set, use the first shard for valid and test
            paths = paths[:1]
        data_path = paths[(epoch - 1) % len(paths)]

        # infer langcode
        src, tgt = self.args.source_lang, self.args.target_lang

        self.datasets[split] = load_langpair_dataset(
            data_path, split, src, self.src_dict, tgt, self.tgt_dict,
            combine=combine, dataset_impl=self.args.dataset_impl,
            upsample_primary=self.args.upsample_primary,
            left_pad_source=self.args.left_pad_source,
            left_pad_target=self.args.left_pad_target,
            max_source_positions=self.args.max_source_positions,
            max_target_positions=self.args.max_target_positions,
            load_alignments=self.args.load_alignments,
            truncate_source=self.args.truncate_source,
            num_buckets=self.args.num_batch_buckets,
            shuffle=(split != 'test'),
            pad_to_multiple=self.args.required_seq_len_multiple,
            prepend_bos=True,
        )

    def build_generator(self, models, args, seq_gen_cls=None, extra_gen_cls_kwargs=None, **unused):
        if self.mode == 'at' or self.mode == 'joint-at' or self.mode == 'joint2-at' or self.mode == 'joint3-at'\
                or self.mode == 'joint-at-only' or self.mode == 'joint2-at-only' or self.mode == 'joint3-at-only':
            if getattr(args, "score_reference", False):
                from fairseq.sequence_scorer import SequenceScorer

                return SequenceScorer(
                    self.target_dictionary,
                    compute_alignment=getattr(args, "print_alignment", False),
                )

            from fairseq.sequence_generator import (
                SequenceGenerator,
                SequenceGeneratorWithAlignment,
            )

            # Choose search strategy. Defaults to Beam Search.
            sampling = getattr(args, "sampling", False)
            sampling_topk = getattr(args, "sampling_topk", -1)
            sampling_topp = getattr(args, "sampling_topp", -1.0)
            diverse_beam_groups = getattr(args, "diverse_beam_groups", -1)
            diverse_beam_strength = getattr(args, "diverse_beam_strength", 0.5)
            match_source_len = getattr(args, "match_source_len", False)
            diversity_rate = getattr(args, "diversity_rate", -1)
            constrained = getattr(args, "constraints", False)
            if (
                    sum(
                        int(cond)
                        for cond in [
                            sampling,
                            diverse_beam_groups > 0,
                            match_source_len,
                            diversity_rate > 0,
                        ]
                    )
                    > 1
            ):
                raise ValueError("Provided Search parameters are mutually exclusive.")
            assert sampling_topk < 0 or sampling, "--sampling-topk requires --sampling"
            assert sampling_topp < 0 or sampling, "--sampling-topp requires --sampling"

            if sampling:
                search_strategy = search.Sampling(
                    self.target_dictionary, sampling_topk, sampling_topp
                )
            elif diverse_beam_groups > 0:
                search_strategy = search.DiverseBeamSearch(
                    self.target_dictionary, diverse_beam_groups, diverse_beam_strength
                )
            elif match_source_len:
                # this is useful for tagging applications where the output
                # length should match the input length, so we hardcode the
                # length constraints for simplicity
                search_strategy = search.LengthConstrainedBeamSearch(
                    self.target_dictionary,
                    min_len_a=1,
                    min_len_b=0,
                    max_len_a=1,
                    max_len_b=0,
                )
            elif diversity_rate > -1:
                search_strategy = search.DiverseSiblingsSearch(
                    self.target_dictionary, diversity_rate
                )
            elif constrained:
                search_strategy = search.LexicallyConstrainedBeamSearch(
                    self.target_dictionary, args.constraints
                )
            else:
                search_strategy = search.BeamSearch(self.target_dictionary)

            if seq_gen_cls is None:
                if getattr(args, "print_alignment", False):
                    seq_gen_cls = SequenceGeneratorWithAlignment
                else:
                    seq_gen_cls = SequenceGenerator
            extra_gen_cls_kwargs = extra_gen_cls_kwargs or {}
            return seq_gen_cls(
                models,
                self.target_dictionary,
                beam_size=getattr(args, "beam", 5),
                max_len_a=getattr(args, "max_len_a", 0),
                max_len_b=getattr(args, "max_len_b", 200),
                min_len=getattr(args, "min_len", 1),
                normalize_scores=(not getattr(args, "unnormalized", False)),
                len_penalty=getattr(args, "lenpen", 1),
                unk_penalty=getattr(args, "unkpen", 0),
                temperature=getattr(args, "temperature", 1.0),
                match_source_len=getattr(args, "match_source_len", False),
                no_repeat_ngram_size=getattr(args, "no_repeat_ngram_size", 0),
                search_strategy=search_strategy,
                **extra_gen_cls_kwargs,
            )
        elif self.mode == 'nat' or self.mode == 'joint-nat' or self.mode == 'joint2-nat' or self.mode == 'joint3-nat'\
                or self.mode == 'joint-nat-only' or self.mode == 'joint2-nat-only' or self.mode == 'joint3-nat-only':
            # add models input to match the API for SequenceGenerator
            from fairseq.iterative_refinement_generator import IterativeRefinementGenerator
            return IterativeRefinementGenerator(
                self.target_dictionary,
                eos_penalty=getattr(args, 'iter_decode_eos_penalty', 0.0),
                max_iter=getattr(args, 'iter_decode_max_iter', 10),
                beam_size=getattr(args, 'iter_decode_with_beam', 1),
                reranking=getattr(args, 'iter_decode_with_external_reranker', False),
                decoding_format=getattr(args, 'decoding_format', None),
                adaptive=not getattr(args, 'iter_decode_force_max_iter', False),
                retain_history=getattr(args, 'retain_iter_history', False))
        else:
            raise ValueError("Mode Error!!!")

    def build_dataset_for_inference(self, src_tokens, src_lengths, constraints=None):
        return LanguagePairDataset(src_tokens, src_lengths, self.source_dictionary,
                                   tgt_dict=self.target_dictionary,
                                   constraints=constraints,
                                   append_bos=True)

    def build_model(self, args):
        model = super().build_model(args)
        if getattr(args, 'eval_bleu', False):
            assert getattr(args, 'eval_bleu_detok', None) is not None, (
                '--eval-bleu-detok is required if using --eval-bleu; '
                'try --eval-bleu-detok=moses (or --eval-bleu-detok=space '
                'to disable detokenization, e.g., when using sentencepiece)'
            )
            detok_args = json.loads(getattr(args, 'eval_bleu_detok_args', '{}') or '{}')
            self.tokenizer = encoders.build_tokenizer(Namespace(
                tokenizer=getattr(args, 'eval_bleu_detok', None),
                **detok_args
            ))

            gen_args = json.loads(getattr(args, 'eval_bleu_args', '{}') or '{}')
            self.sequence_generator = self.build_generator([model], Namespace(**gen_args))
        return model

    def joint_at_nat_train(self, model, criterion, sample, ignore_grad=False):
        # AT
        with torch.autograd.profiler.record_function("forward"):
            loss_at, sample_size, at_logging_output = criterion(model, sample, mode='at')
        # NAT
        sample['prev_target'] = self.inject_noise(sample['target'], noise_type="random_mask")
        loss_nat, sample_size, nat_logging_output = criterion(model, sample, mode='nat', nat_reduction='sum')

        Lambda = model.base_args.lambda_at
        loss = Lambda * loss_at + (1 - Lambda) * loss_nat
        if ignore_grad:
            loss *= 0
        logging_output = {}
        for k, v in nat_logging_output.items():
            if k in at_logging_output.keys():
                logging_output[k] = Lambda * v + (1 - Lambda) * at_logging_output[k]
            else:
                logging_output['nat-' + str(k)] = nat_logging_output[k]

        return loss, sample_size, logging_output

    def train_step(self,
                   sample,
                   model,
                   criterion,
                   optimizer,
                   update_num,
                   ignore_grad=False):
        model.train()
        model.set_num_updates(update_num)
        if self.mode == 'at':
            with torch.autograd.profiler.record_function("forward"):
                loss, sample_size, logging_output = criterion(model, sample)
            if ignore_grad:
                loss *= 0
        elif self.mode == 'nat':
            sample['prev_target'] = self.inject_noise(sample['target'])
            loss, sample_size, logging_output = criterion(model, sample)
            if ignore_grad:
                loss *= 0
        elif self.mode == 'joint-at' or self.mode == 'joint-nat':
            loss, sample_size, logging_output = self.joint_at_nat_train(model, criterion, sample, ignore_grad=ignore_grad)
        elif self.mode == 'joint3-at-only':
            if update_num < self.args.switch_mode_num_update:
                loss, sample_size, logging_output = self.joint_at_nat_train(model, criterion, sample, ignore_grad=ignore_grad)
            if update_num == self.args.switch_mode_num_update:
                print("Enter stage2 training mode")
            with torch.autograd.profiler.record_function("forward"):
                # AT
                net_output = model(**sample['net_input'])
                at_logits = net_output[0]
                at_probs = torch.softmax(at_logits, dim=-1)
                lprobs = torch.log(at_probs)
                with torch.no_grad():
                    model.eval()
                    input_indicator = fetch_nat_input_indicator(at_probs, sample, self.args.strategy, threshold=self.args.threshold)
                    # NAT
                    src_tokens, src_lengths = (
                        sample["net_input"]["src_tokens"],
                        sample["net_input"]["src_lengths"],
                    )

                    init_nat_probs, init_word_ins_mask = self.mask_then_predict(model, sample, src_tokens, src_lengths, input_indicator)

                    word_ins_mask = init_word_ins_mask
                    nat_probs = init_nat_probs
                    nat_lprobs = torch.log(nat_probs)

                    if self.args.kl_strategy == "anneal":
                        factor = (self.args.kl_loss_coef * (1 - (update_num - self.args.switch_mode_num_update) / (
                                self.args.max_update - self.args.switch_mode_num_update))) * word_ins_mask
                    else:
                        raise NotImplementedError

                model.at_transformer.train()

                target = sample['target']
                epsilon = criterion.at_criterion.eps
                ignore_index = criterion.at_criterion.padding_idx
                at_cmlm_kl_loss = (nat_probs * (nat_lprobs - lprobs)).sum(dim=-1)
                if target.dim() == lprobs.dim() - 1:
                    target = target.unsqueeze(-1)
                nll_loss = -lprobs.gather(dim=-1, index=target)
                smooth_loss = -lprobs.sum(dim=-1, keepdim=True)
                if ignore_index is not None:
                    pad_mask = target.eq(ignore_index)
                    nll_loss.masked_fill_(pad_mask, 0.)
                    smooth_loss.masked_fill_(pad_mask, 0.)
                nll_loss = nll_loss.squeeze(-1)
                smooth_loss = smooth_loss.squeeze(-1)
                eps_i = epsilon / lprobs.size(-1)
                loss_at = (1. - epsilon) * nll_loss + eps_i * smooth_loss
                loss = (1. - factor) * loss_at + factor * at_cmlm_kl_loss
                nll_loss = nll_loss.sum()
                loss = loss.sum()
                sample_size = sample['ntokens']
                logging_output = {
                    'loss': loss.data,
                    'nll_loss': nll_loss.data,
                    'ntokens': sample['ntokens'],
                    'nsentences': sample['target'].size(0),
                    'sample_size': sample_size,
                }
        else:
            raise NotImplementedError("Mode Error!!!")

        with torch.autograd.profiler.record_function("backward"):
            optimizer.backward(loss)

        return loss, sample_size, logging_output

    def valid_step(self, sample, model, criterion):
        if self.mode == 'at' or self.mode == 'joint-at' or self.mode == 'joint2-at' or self.mode == 'joint3-at'\
                or self.mode == 'joint-at-only' or self.mode == 'joint2-at-only' or self.mode == 'joint3-at-only':
            model.eval()
            with torch.no_grad():
                loss, sample_size, logging_output = criterion(model, sample, mode='at')
            if self.args.eval_bleu:
                bleu = self._inference_with_bleu(self.sequence_generator, sample, model)
                logging_output['_bleu_sys_len'] = bleu.sys_len
                logging_output['_bleu_ref_len'] = bleu.ref_len
                # we split counts into separate entries so that they can be
                # summed efficiently across workers using fast-stat-sync
                assert len(bleu.counts) == EVAL_BLEU_ORDER
                for i in range(EVAL_BLEU_ORDER):
                    logging_output['_bleu_counts_' + str(i)] = bleu.counts[i]
                    logging_output['_bleu_totals_' + str(i)] = bleu.totals[i]
            return loss, sample_size, logging_output
        elif self.mode == 'nat' or self.mode == 'joint-nat' or self.mode == 'joint2-nat' or self.mode == 'joint3-nat'\
                or self.mode == 'joint-nat-only' or self.mode == 'joint2-nat-only' or self.mode == 'joint3-nat-only':
            model.eval()
            with torch.no_grad():
                sample['prev_target'] = self.inject_noise(sample['target'])
                loss, sample_size, logging_output = criterion(model, sample, mode='nat')
                if self.args.eval_bleu:
                    bleu = self._inference_with_bleu(self.sequence_generator, sample, model)
                    logging_output['_bleu_sys_len'] = bleu.sys_len
                    logging_output['_bleu_ref_len'] = bleu.ref_len
                    # we split counts into separate entries so that they can be
                    # summed efficiently across workers using fast-stat-sync
                    assert len(bleu.counts) == EVAL_BLEU_ORDER
                    for i in range(EVAL_BLEU_ORDER):
                        logging_output['_bleu_counts_' + str(i)] = bleu.counts[i]
                        logging_output['_bleu_totals_' + str(i)] = bleu.totals[i]
            return loss, sample_size, logging_output
        else:
            raise ValueError("Mode Error")

    def reduce_metrics(self, logging_outputs, criterion):
        super().reduce_metrics(logging_outputs, criterion)
        if self.args.eval_bleu:

            def sum_logs(key):
                return sum(log.get(key, 0) for log in logging_outputs)

            counts, totals = [], []
            for i in range(EVAL_BLEU_ORDER):
                counts.append(sum_logs('_bleu_counts_' + str(i)))
                totals.append(sum_logs('_bleu_totals_' + str(i)))

            if max(totals) > 0:
                # log counts as numpy arrays -- log_scalar will sum them correctly
                metrics.log_scalar('_bleu_counts', np.array(counts))
                metrics.log_scalar('_bleu_totals', np.array(totals))
                metrics.log_scalar('_bleu_sys_len', sum_logs('_bleu_sys_len'))
                metrics.log_scalar('_bleu_ref_len', sum_logs('_bleu_ref_len'))

                def compute_bleu(meters):
                    import inspect
                    import sacrebleu
                    fn_sig = inspect.getfullargspec(sacrebleu.compute_bleu)[0]
                    if 'smooth_method' in fn_sig:
                        smooth = {'smooth_method': 'exp'}
                    else:
                        smooth = {'smooth': 'exp'}
                    bleu = sacrebleu.compute_bleu(
                        correct=meters['_bleu_counts'].sum,
                        total=meters['_bleu_totals'].sum,
                        sys_len=meters['_bleu_sys_len'].sum,
                        ref_len=meters['_bleu_ref_len'].sum,
                        **smooth
                    )
                    return round(bleu.score, 2)

                metrics.log_derived('bleu', compute_bleu)

    def max_positions(self):
        """Return the max sentence length allowed by the task."""
        return (self.args.max_source_positions, self.args.max_target_positions)

    @property
    def source_dictionary(self):
        """Return the source :class:`~fairseq.data.Dictionary`."""
        return self.src_dict

    @property
    def target_dictionary(self):
        """Return the target :class:`~fairseq.data.Dictionary`."""
        return self.tgt_dict

    def _inference_with_bleu(self, generator, sample, model):
        import sacrebleu

        def decode(toks, escape_unk=False):
            s = self.tgt_dict.string(
                toks.int().cpu(),
                self.args.eval_bleu_remove_bpe,
                # The default unknown string in fairseq is `<unk>`, but
                # this is tokenized by sacrebleu as `< unk >`, inflating
                # BLEU scores. Instead, we use a somewhat more verbose
                # alternative that is unlikely to appear in the real
                # reference, but doesn't get split into multiple tokens.
                unk_string=(
                    "UNKNOWNTOKENINREF" if escape_unk else "UNKNOWNTOKENINHYP"
                ),
            )
            if self.tokenizer:
                s = self.tokenizer.decode(s)
            return s

        gen_out = self.inference_step(generator, [model], sample, prefix_tokens=None)
        hyps, refs = [], []
        for i in range(len(gen_out)):
            hyps.append(decode(gen_out[i][0]['tokens']))
            refs.append(decode(
                utils.strip_pad(sample['target'][i], self.tgt_dict.pad()),
                escape_unk=True,  # don't count <unk> as matches to the hypo
            ))
        if self.args.eval_bleu_print_samples:
            logger.info('example hypothesis: ' + hyps[0])
            logger.info('example reference: ' + refs[0])
        if self.args.eval_tokenized_bleu:
            return sacrebleu.corpus_bleu(hyps, [refs], tokenize='none')
        else:
            return sacrebleu.corpus_bleu(hyps, [refs])

    def inject_noise(self, target_tokens, noise_type=None, mask_rate=None, input_indicator=None):
        def _random_delete(target_tokens):
            pad = self.tgt_dict.pad()
            bos = self.tgt_dict.bos()
            eos = self.tgt_dict.eos()

            max_len = target_tokens.size(1)
            target_mask = target_tokens.eq(pad)
            target_score = target_tokens.clone().float().uniform_()
            target_score.masked_fill_(
                target_tokens.eq(bos) | target_tokens.eq(eos), 0.0)
            target_score.masked_fill_(target_mask, 1)
            target_score, target_rank = target_score.sort(1)
            target_length = target_mask.size(1) - target_mask.float().sum(
                1, keepdim=True)

            # do not delete <bos> and <eos> (we assign 0 score for them)
            target_cutoff = 2 + ((target_length - 2) * target_score.new_zeros(
                target_score.size(0), 1).uniform_()).long()
            target_cutoff = target_score.sort(1)[1] >= target_cutoff

            prev_target_tokens = target_tokens.gather(1, target_rank).masked_fill_(target_cutoff, pad).gather(
                1, target_rank.masked_fill_(target_cutoff, max_len).sort(1)[1])
            prev_target_tokens = prev_target_tokens[:, :prev_target_tokens.
                                                    ne(pad).sum(1).max()]

            return prev_target_tokens

        def _random_mask(target_tokens):
            pad = self.tgt_dict.pad()
            bos = self.tgt_dict.bos()
            eos = self.tgt_dict.eos()
            unk = self.tgt_dict.unk()

            target_masks = target_tokens.ne(pad) & \
                           target_tokens.ne(bos) & \
                           target_tokens.ne(eos)
            target_score = target_tokens.clone().float().uniform_()
            target_score.masked_fill_(~target_masks, 2.0)
            target_length = target_masks.sum(1).float()
            target_length = target_length * target_length.clone().uniform_()
            target_length = target_length + 1  # make sure to mask at least one token.

            _, target_rank = target_score.sort(1)
            target_cutoff = new_arange(target_rank) < target_length[:, None].long()
            prev_target_tokens = target_tokens.masked_fill(
                target_cutoff.scatter(1, target_rank, target_cutoff), unk)
            return prev_target_tokens

        def _full_mask(target_tokens):
            pad = self.tgt_dict.pad()
            bos = self.tgt_dict.bos()
            eos = self.tgt_dict.eos()
            unk = self.tgt_dict.unk()

            target_mask = target_tokens.eq(bos) | target_tokens.eq(
                eos) | target_tokens.eq(pad)
            return target_tokens.masked_fill(~target_mask, unk)

        def _mask_by_rate(target_tokens, mask_rate):
            pad = self.tgt_dict.pad()
            bos = self.tgt_dict.bos()
            eos = self.tgt_dict.eos()
            unk = self.tgt_dict.unk()

            target_masks = target_tokens.ne(pad) & \
                           target_tokens.ne(bos) & \
                           target_tokens.ne(eos)
            target_score = target_tokens.clone().float().uniform_()
            target_score.masked_fill_(~target_masks, 2.0)
            target_length = target_masks.sum(1).float()
            target_length = target_length * target_length.clone().uniform_()
            target_length = target_length + 1  # make sure to mask at least one token.

            _, target_rank = target_score.sort(1)
            target_cutoff = new_arange(target_rank) < target_length[:, None].long()
            prev_target_tokens = target_tokens.masked_fill(
                target_cutoff.scatter(1, target_rank, target_cutoff), unk)
            return prev_target_tokens

        def _non_match_mask(target_tokens, input_indicator):
            pad = self.tgt_dict.pad()
            bos = self.tgt_dict.bos()
            eos = self.tgt_dict.eos()
            unk = self.tgt_dict.unk()

            target_mask = target_tokens.eq(bos) | target_tokens.eq(
                eos) | target_tokens.eq(pad) | input_indicator
            return target_tokens.masked_fill(~target_mask, unk)

        if noise_type is None:
            if self.args.noise == 'random_delete':
                return _random_delete(target_tokens)
            elif self.args.noise == 'random_mask':
                return _random_mask(target_tokens)
            elif self.args.noise == 'full_mask':
                return _full_mask(target_tokens)
            elif self.args.noise == 'non_match_mask':
                return _non_match_mask(target_tokens, input_indicator)
            elif self.args.noise == 'no_noise':
                return target_tokens
            else:
                raise NotImplementedError
        else:
            if noise_type == 'random_delete':
                return _random_delete(target_tokens)
            elif noise_type == 'random_mask':
                return _random_mask(target_tokens)
            elif noise_type == 'full_mask':
                return _full_mask(target_tokens)
            elif noise_type == 'non_match_mask':
                return _non_match_mask(target_tokens, input_indicator)
            elif noise_type == 'no_noise':
                return target_tokens
            else:
                raise NotImplementedError

    def set_mode(self, mode):
        self.mode = mode

    def mask_then_predict(self, model, sample, src_tokens, src_lengths, input_indicator):

        sample['prev_target'] = self.inject_noise(sample['target'], noise_type="non_match_mask", input_indicator=input_indicator)
        tgt_tokens, prev_output_tokens = sample["target"], sample["prev_target"]

        outputs = model.nat_transformer(src_tokens, src_lengths, prev_output_tokens, tgt_tokens=tgt_tokens)
        # word_ins_out(before normalization): -> (B, L, V); word_ins_tgt -> (B, L); word_ins_mask -> (B, L)
        word_ins_out, word_ins_tgt, init_word_ins_mask = outputs["word_ins"]['out'], outputs["word_ins"]['tgt'], \
                                                         outputs["word_ins"]['mask']
        init_nat_probs = torch.softmax(word_ins_out, dim=-1)
        # nat_output_states = outputs["word_ins"]["inner_states"]

        return init_nat_probs, init_word_ins_mask