import math

import torch
import torch.nn.functional as F

from fairseq.criterions.sentence_prediction import SentencePredictionCriterion
from fairseq.criterions import register_criterion
from fairseq import metrics

@register_criterion('extraction_criterion')
class ExtractionCriterion(SentencePredictionCriterion):
    def __init__(self, task, classification_head_name, 
        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

    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'

        logits, _ = model(
            **sample['net_input'],
            features_only=True,
            classification_head_name=self.classification_head_name,
        )
        
        bsz, max_sent_num, _ = logits.size()

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

        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

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

        return 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)
        nsentences = sum(log.get('nsentences', 0) for log in logging_outputs)
        sample_size = sum(log.get('sample_size', 0) for log in logging_outputs)

        metrics.log_scalar('loss', 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('accuracy', 100.0 * ncorrect / nsentences, nsentences, round=1)
