import torch
import torch.nn as nn

from opencd.registry import MODELS


def bcl_loss(
    pred, 
    target, 
    margin=2.0, 
    eps=1e-4, 
    ignore_index=255, 
    **kwargs):
    pred = pred.squeeze()
    target = target.squeeze()
    assert pred.size() == target.size() and target.numel() > 0
    mask = (target != ignore_index).float()
    target = target * mask
    utarget = 1 - target
    n_u = utarget.sum() + eps
    n_c = target.sum() + eps
    loss = torch.sum(utarget * torch.pow(pred, 2) * mask) / n_u + \
        torch.sum(target * torch.pow(torch.clamp(margin - pred, min=0.), 2)) / n_c
    return loss


@MODELS.register_module()
class BCLLoss(nn.Module):
    """Batch-balanced Contrastive Loss"""

    def __init__(
        self, 
        margin=2.0, 
        loss_weight=1.0, 
        ignore_index=255,
        loss_name='bcl_loss',
        **kwargs):
        super().__init__()
        self.margin = margin
        self.loss_weight = loss_weight
        self.ignore_index = ignore_index
        self._loss_name = loss_name

    def forward(self,
                pred,
                target,
                **kwargs):

        loss = self.loss_weight * bcl_loss(
            pred, target, self.margin, self.ignore_index)
        return loss

    @property
    def loss_name(self):
        """Loss Name.
        This function must be implemented and will return the name of this
        loss function. This name will be used to combine different loss items
        by simple sum operation. In addition, if you want this loss item to be
        included into the backward graph, `loss_` must be the prefix of the
        name.
        Returns:
            str: The name of this loss item.
        """
        return self._loss_name