# 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 math
import numpy as np

import torch
import torch.nn.functional as F

from fairseq import metrics, utils
from fairseq.criterions import FairseqCriterion, register_criterion


@register_criterion('mask_co_leaner')
class MaskLeanerCoLoss(FairseqCriterion):
    """
    Implementation for the loss used in masked language model (MLM) training.
    """

    def __init__(self, args, task):
        super(MaskLeanerCoLoss, self).__init__(args, task)

        self.vocab = self.task.source_dictionary
        self.mask_idx = self.task.mask_idx
        self.mask_prob = self.task.args.mask_prob
        self.leave_unmasked_prob = self.task.args.leave_unmasked_prob
        self.random_token_prob = self.task.args.random_token_prob
        self.rand_or_unmask_prob = self.random_token_prob + self.leave_unmasked_prob

        self.mask_whole_words = self.task.args.mask_whole_words
        self.freq_weighted_replacement = self.task.args.freq_weighted_replacement

        self.masker_lambda = args.masker_lambda
        self.masker_m = args.masker_m
        self.avg_baseline = args.masker_m < 0.0
        self.do_deterministic = args.masker_deterministic
        self.masker_eps = args.masker_eps
        if self.random_token_prob > 0.0:
            if self.freq_weighted_replacement:
                weights = np.array(self.vocab.count)
            else:
                weights = np.ones(len(self.vocab))
            weights[:self.vocab.nspecial] = 0
            self.weights = weights / weights.sum()
            self.register_buffer('random_weights', torch.tensor(self.weights).type(torch.float32))

        self.op_dict={
            'relu': F.relu,
            'identity': lambda x: x,
        }
        self.masker_loss_op = self.op_dict[args.masker_op_name]

    @staticmethod
    def add_args(parser):
        """Add criterion-specific arguments to the parser."""
        super(MaskLeanerCoLoss,
              MaskLeanerCoLoss).add_args(parser)

        parser.add_argument('--masker_lambda', default=0.5, type=float, metavar='D',
                            help='weight for the masker loss')
        parser.add_argument('--masker_m', default=0.5, type=float,
                            help='m for the masker loss. If m < 0, then batch average')
        parser.add_argument('--masker_eps', default=0.2, type=float,
                            help='epsilon used for clipping')
        parser.add_argument('--masker_deterministic', default=False,
                            action='store_true',
                            help='is the mask generated by sampling?')
        parser.add_argument('--masker_op_name', default='identity', type=str,
                            help='\sigma in the master loss')

    def forward(self, model, sample, reduce=True):
        """Compute the loss for the given sample.

        Returns a tuple with three elements:
        1) the loss
        2) the sample size, which is used as the denominator for the gradient
        3) logging outputs to display while training
        """
        # compute MLM loss
        # model.learner model.lm

        raw_inps = sample["net_input"]["src_tokens"]
        raw_targets = sample['target']
        raw_masked_tokens = raw_targets.ne(self.padding_idx)
        inps = raw_targets * raw_masked_tokens + \
               raw_inps * (raw_masked_tokens ^ True)
        sz = inps.size(-1) # all batches should be the same length
        num_mask = int(sz * 0.15)

        masker_out = model.masker(inps)[0]#.view(inps.size(0), -1)
        with torch.no_grad():
            adz = torch.ones_like(masker_out) * 1e-5
            log_masking_softmax = torch.log(masker_out + adz)
            p_logp = log_masking_softmax * masker_out

            masker_entropy = torch.sum(p_logp) * -1.0

            token_length = masker_out.size(1)

            index5 = 5 if token_length > 4 else token_length
            index2 = 2 if token_length > 2 else token_length

            top5_masker, _ = torch.topk(masker_out, index5, dim=-1)
            top2_dist = top5_masker[:, 0] - top5_masker[:, index2-1]
            top5_dist = top5_masker[:, 0] - top5_masker[:, index5-1]
            top2_dist = torch.sum(top2_dist)
            top5_dist = torch.sum(top5_dist)

            #print(masker_entropy, inps.shape)

        #print('masker 1 shape', masker_out.shape)

        if num_mask == 0:
            num_mask = 1

        if self.do_deterministic:
            masked_tokens, masked_idxes = torch.topk(masker_out.float(), num_mask, dim=-1)
        else:
            with torch.no_grad():
                #t_masker_out = torch.clamp(masker_out * float(num_mask), 0, 1)
                #random_s = torch.bernoulli(t_masker_out).type(torch.bool)
                #masked_idxes = random_s  # not not index, but table of True of False##
                # print ('masker')
                masked_idxes = torch.multinomial(masker_out.float(), num_mask, replacement=False)



        with torch.no_grad():

            labels = torch.full_like(inps, self.padding_idx)

            raw_masked_pos = torch.full_like(inps, False).type(torch.bool)

            x_idx = torch.arange(end=inps.size(0), device=inps.device).unsqueeze(-1)
            x_idx = x_idx.expand_as(masked_idxes)

            #from IPython import embed
            #embed()
            raw_masked_pos[(x_idx, masked_idxes)] = True

            #raw_masked_pos.scatter_(-1, masked_idxes, True)

            raw_masked_pos[inps == self.padding_idx] = False

            #labels[(x_idx, masked_idxes)] = inps[(x_idx, masked_idxes)]
            labels[raw_masked_pos] = inps[raw_masked_pos]

            new_masked_inps = inps * raw_masked_pos.logical_not() + \
                              torch.full_like(inps, self.mask_idx) * raw_masked_pos

            probs = torch.ones_like(inps, dtype=torch.float)

            probs[raw_masked_pos] = torch.rand_like(inps, dtype=torch.float)[raw_masked_pos]

            non_masked_pos = probs < self.leave_unmasked_prob

            new_masked_inps[non_masked_pos] = inps[non_masked_pos]

            rand_replace_pos = (probs >= self.leave_unmasked_prob) & (probs < self.rand_or_unmask_prob)

            num_gen = rand_replace_pos.long().sum()

            #real_mask_pos = raw_masked_pos & (non_masked_pos.logical_not()) & (rand_replace_pos.logical_not())
            #real_from_raw = real_mask_pos[raw_masked_pos]

            if num_gen > 0:
                rand_token = torch.multinomial(self.random_weights.float(), num_gen, replacement=True)

                rand_token.type(inps.type())
                new_masked_inps[rand_replace_pos] = rand_token

            new_inp = new_masked_inps

            if model.training:
                sample['target'] = labels
                sample['net_input']["src_tokens"] = new_inp
                
        masked_tokens = sample['target'].ne(self.padding_idx)
        sample_size = masked_tokens.int().sum().item()

        raw_masked_pos = sample['target'].ne(self.padding_idx)
        real_mask_pos = sample['net_input']["src_tokens"].eq(self.mask_idx)
        real_from_raw = real_mask_pos[raw_masked_pos]

        # (Rare case) When all tokens are masked, the model results in empty
        # tensor and gives CUDA error.
        if sample_size == 0:
            masked_tokens = None

        logits = model(**sample['net_input'], masked_tokens=masked_tokens)[0]
        targets = model.get_targets(sample, [logits])

        batch_sz, len_token = sample['net_input']["src_tokens"].size()

        if sample_size != 0:
            targets = targets[masked_tokens]

        loss_ = F.nll_loss( # [sample_size]
            F.log_softmax(
                logits.view(-1, logits.size(-1)),
                dim=-1,
                dtype=torch.float32,
            ),
            targets.view(-1),
            reduction='none',
            ignore_index=self.padding_idx,
        )

        loss, weight_mean = None, None
        raw_masker_out = masker_out
        masker_out = masker_out[masked_tokens]
        if sample_size % batch_sz == 0:
            loss_ = loss_.view(batch_sz, -1)
            masker_out = masker_out.view(batch_sz, -1)

            #
            #weight = torch.prod( (masker_out.detach()) * (token_length * 1.0),
            #                     dim=-1, keepdim=True)
            weight = torch.log((masker_out.detach()) * (token_length * 1.0)) * -1.0
            weight = torch.sum(weight, dim=-1, keepdim=True)
            weight = torch.exp(weight)

            #weight = weight.reciprocal()

            #print(weight.shape, torch.mean(weight).item(), torch.max(weight).item(), torch.min(weight).item())
            # put the average weight into logs

            weight = torch.clamp(weight, 1.0 - self.masker_eps, 1.0 + self.masker_eps)
            with torch.no_grad():
                weight_mean = torch.sum(weight) # should be smaller than 1.0

            weight = weight / torch.mean(weight) # normalize

            # print (weight.size(), loss_.size())
            loss_re = weight * loss_  # sample size
            loss = torch.sum(loss_re)

            masker_out = masker_out.view(-1)
        else:
            mask_per_sent = torch.sum(masked_tokens.int(), 1)
            masker_out1 = masker_out.detach()

            loss_b, weight, idx = [], [], 0
            for num in mask_per_sent:
                c, p = torch.zeros((1,), device=masker_out1.device), torch.zeros((1,), device=masker_out1.device)
                for _ in range(num.item()):
                    c += loss_[idx]
                    p  = p + torch.log( (token_length * 1.0) * masker_out1[idx])
                    idx += 1
                loss_b.append(c)
                weight.append(torch.exp(p*-1.0))

            # numpy array of torch tensor!!
            loss_b, weight = np.array(loss_b), np.array(weight)

            weight = np.clip(weight, 1.0 - self.masker_eps, 1.0 + self.masker_eps)
            weight_mean = np.sum(weight)
            weight_mean = torch.tensor(weight_mean)

            a = np.sum(weight)
            normalize_term = a / sample['nsentences']

            # print(a, normalize_term)
            if torch.is_tensor(normalize_term):
                normalize_term = normalize_term.item()
            weight = weight / normalize_term
            
            loss_re = loss_b * weight
            loss = np.sum(loss_re)
            # print (loss)

        #import IPython
        #IPython.embed()

        #import IPython
        #IPython.embed()
        #print(pred_softmax.shape, targets.shape)
        #print(real_from_raw.shape, loss_.shape, masker_out.shape)
        bert_loss = loss_.detach().view(-1)[real_from_raw]
        masker_out = masker_out[real_from_raw]

        with torch.no_grad():
            # put batch mean into logging!
            batch_mean = bert_loss.mean(dim=-1).sum()

        if self.avg_baseline:
            baseline = bert_loss.mean()
        else:
            baseline = torch.ones_like(bert_loss) * self.masker_m
            baseline = torch.log(baseline) * -1.0
        bert_loss = bert_loss - baseline
        #masker_loss = bert_loss * masker_out * -1.0
        avoid_div_zero = torch.ones_like(masker_out) * 1e-4

        bert_loss = self.masker_loss_op(bert_loss)
        masker_loss = torch.log(masker_out + avoid_div_zero) * bert_loss * -1.0
        masker_loss = masker_loss.sum()

        total_loss = masker_loss * self.masker_lambda + loss



        logging_output = {
            'loss': utils.item(loss.data) if reduce else loss.data,
            'ntokens': sample['ntokens'],
            'nsentences': sample['nsentences'],
            'sample_size': sample_size,
            'masker_loss':utils.item(masker_loss.data) if reduce else masker_loss.data,
            'total_loss': utils.item(total_loss.data) if reduce else total_loss.data,
            'masker_entropy': utils.item(masker_entropy.data) if reduce else masker_entropy.data,
            'top2_dist': utils.item(top2_dist.data) if reduce else top2_dist.data,
            'top5_dist': utils.item(top5_dist.data) if reduce else top5_dist.data,
            'weight_mean': utils.item(weight_mean.data) if reduce else weight_mean.data,
            'batch_mean': utils.item(batch_mean.data) if reduce else batch_mean.data,
        }
        return total_loss, sample_size, logging_output

    @staticmethod
    def reduce_metrics(logging_outputs) -> None:
        """Aggregate logging outputs from data parallel training."""
        loss_sum = sum(log.get('loss', 0) for log in logging_outputs)
        masker_loss_sum = sum(log.get('masker_loss', 0) for log in logging_outputs)
        total_loss_sum = sum(log.get('total_loss', 0) for log in logging_outputs)
        masker_entropy = sum(log.get('masker_entropy', 0) for log in logging_outputs)
        top2_dist = sum(log.get('top2_dist', 0) for log in logging_outputs) / len(logging_outputs)
        top5_dist = sum(log.get('top5_dist', 0) for log in logging_outputs) / len(logging_outputs)
        weight_mean = sum(log.get('weight_mean', 0) for log in logging_outputs) / len(logging_outputs)
        batch_mean = sum(log.get('batch_mean', 0) for log in logging_outputs) / len(logging_outputs)
        #print(len(logging_outputs), 'len logging outputs')

        sample_size = sum(log.get('sample_size', 0) for log in logging_outputs)
        nsentences = sum(log.get('nsentences', 0) for log in logging_outputs)

        metrics.log_scalar('loss', loss_sum / sample_size / math.log(2), sample_size, round=3)
        metrics.log_scalar('masker_entropy', masker_entropy / nsentences , 8, round=3)

        metrics.log_scalar('top2_dist', top2_dist / nsentences, 32, round=4)
        metrics.log_scalar('top5_dist', top5_dist / nsentences, 32, round=4)
        metrics.log_scalar('weight_mean', weight_mean / nsentences, 32, round=2)
        metrics.log_scalar('batch_mean', batch_mean / nsentences, 32, round=2)
        #metrics.log_scalar('nsentences', nsentences, 1, round=1)
        metrics.log_scalar('masker_loss', masker_loss_sum / sample_size / math.log(2)  , sample_size, round=5)
        metrics.log_scalar('total_loss', total_loss_sum / sample_size / math.log(2), sample_size, round=3)
        metrics.log_derived('ppl', lambda meters: round(2**meters['loss'].avg, 3))

    @staticmethod
    def logging_outputs_can_be_summed() -> bool:
        """
        Whether the logging outputs returned by `forward` can be summed
        across workers prior to calling `reduce_metrics`. Setting this
        to True will improves distributed training speed.
        """
        return True
