import torch
import torch.nn as nn
import torch.nn.functional as F
from uniperceiver.config import configurable
from .build import LOSSES_REGISTRY




@LOSSES_REGISTRY.register()
class LabelSmoothingCrossEntropy(nn.Module):
    @configurable
    def __init__(
        self,
        *,
        label_smoothing,
        loss_weight,
        loss_fp32,
    ):
        super(LabelSmoothingCrossEntropy, self).__init__()
        self.label_smoothing = label_smoothing
        self.confidence = 1.0 - self.label_smoothing
        if not isinstance(loss_weight, float):
            self.loss_weight = 1.0
        else:
            self.loss_weight = loss_weight
        self.loss_fp32 = loss_fp32

    @classmethod
    def from_config(cls, cfg):
        return {
            "label_smoothing": cfg.LOSSES.LABELSMOOTHING,
            'loss_weight': getattr(cfg.LOSSES, 'LOSS_WEIGHT', None),
            'loss_fp32': getattr(cfg.LOSSES, 'LOSS_FP32', False),
        }

    def Forward(self, x, target):
        if self.loss_fp32 and x.dtype != torch.float32:
            logprobs = F.log_softmax(x, dim=-1,
                                     dtype=torch.float32).to(x.dtype)
        else:
            logprobs = F.log_softmax(x, dim=-1)
        nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1))
        nll_loss = nll_loss.squeeze(1)
        smooth_loss = -logprobs.mean(dim=-1)
        loss = self.confidence * nll_loss + self.label_smoothing * smooth_loss
        return loss.mean()

    def forward(self, outputs_dict):
        ret = {}

        for logit, target, loss_identification in zip(outputs_dict['logits'],
                                            outputs_dict['targets'],
                                            outputs_dict['loss_names']):


            loss = self.Forward(logit, target)
            if self.loss_weight != 1.0:
                loss *= self.loss_weight
            loss_name = 'LabelSmoothing'
            if len(loss_identification) > 0:
                loss_name = loss_name + f' ({loss_identification})'
            ret.update({loss_name: loss})


        return ret