import math

import torch
import torch.nn.functional as F

from torch.nn import CrossEntropyLoss
from fairseq.criterions.sentence_prediction import SentencePredictionCriterion
from fairseq.criterions import register_criterion
from fairseq import metrics

@register_criterion('extraction_criterion_w_lang_classifier')
class ExtractionCriterionWLangClassifier(SentencePredictionCriterion):
    def __init__(self, task, classification_head_name, lang_loss_weight=1,
        regression_target=False, ignore_index=2
    ):
        super().__init__(task, classification_head_name, regression_target)
        self.loss = torch.nn.CrossEntropyLoss(reduction='sum', ignore_index=ignore_index)
        self.classification_head_name = task.args.extraction_head_name
        self.lang_loss = CrossEntropyLoss(reduction='sum')
        self.langs = task.args.summ_langs + task.args.denoise_langs
        self.lang_loss_weight = lang_loss_weight

    @staticmethod
    def add_args(parser):
        # fmt: off
        parser.add_argument('--classification-head-name',
                            default='sentence_classification_head',
                            help='name of the classification head to use')
        # fmt: on
        parser.add_argument('--lang-loss-weight', default=1.0, type=float)

    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
        """
        assert (
            hasattr(model, 'classification_heads')
            and self.classification_head_name in model.classification_heads
        ), 'model must provide sentence classification head for --criterion=sentence_prediction'

        net_output = model(
            **sample['net_input'],
            features_only=True,
            classification_head_name=self.classification_head_name,
        )

        logits = net_output[0]
        bsz, max_sent_num, _ = logits.size()

        targets = sample['ext_label']
        ext_mask = sample['ext_mask']
        sample_size = bsz
        sample_sent_num = ext_mask.numel()
        ext_sent_num = (~ext_mask).sum()

        if not self.regression_target:
            logits = logits.view(bsz * max_sent_num, -1).float()
            targets = targets.view(-1).long()
            loss = self.loss(logits, targets)
        else:
            raise NotImplementedError

        lang_classifier_outs = net_output[1]['lang_cls_out']
        lang_labels = sample['lang_label']

        lang_loss = self.lang_loss(
            lang_classifier_outs, 
            lang_labels
        )

        sum_loss = loss + lang_loss * self.lang_loss_weight
        logging_output = {
            'loss': sum_loss.data,
            'lang_loss': lang_loss.data,
            'ext_loss': loss.data,
            'ntokens': sample['ntokens'],
            'nsentences': sample_sent_num,
            'sample_size': bsz,
            'n_ext_sent': ext_sent_num
        }
        if not self.regression_target:
            preds = logits.argmax(dim=-1)
            ncorrect = (preds == targets) * (~ext_mask.view(-1))
            logging_output['ncorrect'] = ncorrect.sum()

            lang_preds = lang_classifier_outs.argmax(dim=-1)
            lang_ncorrect = (lang_preds == lang_labels)
            logging_output['lang_ncorrect'] = lang_ncorrect.sum()

        return sum_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)
        lang_loss_sum = sum(log.get('lang_loss', 0) for log in logging_outputs)
        ext_loss_sum = sum(log.get('ext_loss', 0) for log in logging_outputs)
        nsentences = sum(log.get('nsentences', 0) for log in logging_outputs)
        sample_size = sum(log.get('sample_size', 0) for log in logging_outputs)
        n_ext_sent = sum(log.get('n_ext_sent', 0) for log in logging_outputs)
        n_ext_sent += 1

        metrics.log_scalar('loss', loss_sum / sample_size / math.log(2), sample_size, round=3)
        metrics.log_scalar('lang_loss', lang_loss_sum / sample_size / math.log(2), sample_size, round=3)
        metrics.log_scalar('ext_loss', ext_loss_sum / sample_size / math.log(2), sample_size, round=3)

        if len(logging_outputs) > 0 and 'ncorrect' in logging_outputs[0]:
            ncorrect = sum(log.get('ncorrect', 0) for log in logging_outputs)
            metrics.log_scalar('acc', 100.0 * ncorrect / n_ext_sent, n_ext_sent, round=1)

            lang_ncorrect = sum(log.get('lang_ncorrect', 0) for log in logging_outputs)
            metrics.log_scalar('lang_acc', 100.0 * lang_ncorrect / sample_size, sample_size, round=1)
