import torch.nn as nn
import torch.nn.functional as F
import torch


class FocalLoss(nn.Module):
    def __init__(self, alpha=1, gamma=0, size_average=True, ignore_index=255):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.ignore_index = ignore_index
        self.size_average = size_average

    def forward(self, inputs, targets):
        ce_loss = F.cross_entropy(
            inputs, targets, reduction='none', ignore_index=self.ignore_index)
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
        if self.size_average:
            return focal_loss.mean()
        else:
            return focal_loss.sum()


class BCEWithLogitsLossWithIgnoreIndex(nn.Module):
    def __init__(self, reduction='mean', ignore_index=255):
        super().__init__()
        self.reduction = reduction
        self.ignore_index = ignore_index

    def forward(self, inputs, targets, weight=None):
        # inputs of size B x C x H x W
        n_cl = torch.tensor(inputs.shape[1]).to(inputs.device)
        labels_new = torch.where(targets != self.ignore_index, targets, n_cl)
        # replace ignore with numclasses + 1 (to enable one hot and then remove it)
        targets = F.one_hot(labels_new, inputs.shape[1] + 1).float().permute(0, 3, 1, 2)
        targets = targets[:, :inputs.shape[1], :, :]  # remove 255 from 1hot
        # targets is B x C x H x W so shape[1] is C

        loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')

        # loss has shape B x C x H x W
        loss = loss.sum(dim=1)  # sum the contributions of the classes

        if weight is not None:
            loss = loss * weight

        if self.reduction == 'mean':
            # return loss.mean()
            # if targets have only zeros, we skip them
            return torch.masked_select(loss, targets.sum(dim=1) != 0).mean()
        elif self.reduction == 'sum':
            # return loss.sum()
            return torch.masked_select(loss, targets.sum(dim=1) != 0).sum()
        else:
            # return loss
            return loss * targets.sum(dim=1)


class UnseenAugLoss(nn.Module):
    def __init__(self):
        super(UnseenAugLoss, self).__init__()
        self.CE = nn.CrossEntropyLoss()

    def forward(self, opts, predict):
        """
        :param predict: a unknown mask prediction b n h w
        :return: loss
        """

        x = predict.view(predict.shape[0], predict.shape[1], -1)  # b n hw
        x = F.normalize(x, p=2, dim=2)
        simmat = torch.einsum('bic,bjc->bij', x, x.detach())  # b n n
        logits = -torch.log(torch.softmax(simmat, dim=-1))
        l = [torch.diag(i) for i in logits]
        l = torch.stack(l)
        loss = torch.mean(l)

        return loss
