import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from fastNLP.core.losses import LossBase


reduce_func = {
    'none': lambda x, mask: x*mask,
    'sum': lambda x, mask: (x*mask).sum(),
    'mean': lambda x, mask: (x*mask).sum() / mask.sum(),
}


class LabelSmoothCrossEntropy(nn.Module):
    def __init__(self, smoothing=0.1, ignore_index=-100, reduction='mean'):
        global reduce_func
        super().__init__()
        if smoothing < 0 or smoothing > 1:
            raise ValueError('invalid smoothing value: {}'.format(smoothing))
        self.smoothing = smoothing
        self.ignore_index = ignore_index
        if reduction not in reduce_func:
            raise ValueError('invalid reduce type: {}'.format(reduction))
        self.reduce_func = reduce_func[reduction]

    def forward(self, input, target):
        input = F.log_softmax(input, dim=1)  # [N, C, ...]
        smooth_val = self.smoothing / input.size(1)  # [N, C, ...]
        target_logit = input.new_full(input.size(), fill_value=smooth_val)
        target_logit.scatter_(1, target[:, None], 1 - self.smoothing)
        result = -(target_logit * input).sum(1)  # [N, ...]
        mask = (target != self.ignore_index).float()
        return self.reduce_func(result, mask)


class SmoothCE(LossBase):
    def __init__(self, pred=None, target=None, **kwargs):
        super().__init__()
        self.loss_fn = LabelSmoothCrossEntropy(**kwargs)
        self._init_param_map(pred=pred, target=target)

    def get_loss(self, pred, target):
        return self.loss_fn(pred, target)


if __name__ == '__main__':
    loss_fn = nn.CrossEntropyLoss(ignore_index=0)
    sm_loss_fn = LabelSmoothCrossEntropy(smoothing=0, ignore_index=0)
    predict = torch.tensor([[0, 0.2, 0.7, 0.1, 0],
                            [0, 0.9, 0.2, 0.1, 0],
                            [1, 0.2, 0.7, 0.1, 0]])
    target = torch.tensor([2, 1, 0])
    loss = loss_fn(predict, target)
    sm_loss = sm_loss_fn(predict, target)
    print(loss, sm_loss)