# import torch
# import torch.nn as nn
# import torch.nn.functional as F
# from utils import *
from PepperPepper.environment import torch, nn, F
from PepperPepper.IRSTD.tools.utils import *

class SoftLoULoss(nn.Module):
    def __init__(self):
        super(SoftLoULoss, self).__init__()

    def forward(self, pred, target):
        pred = F.sigmoid(pred)
        smooth = 1

        intersection = pred * target

        intersection_sum = torch.sum(intersection, dim=(1,2,3))
        pred_sum = torch.sum(pred, dim=(1,2,3))
        target_sum = torch.sum(target, dim=(1,2,3))
        loss = (intersection_sum + smooth) / \
               (pred_sum + target_sum - intersection_sum + smooth)

        loss = 1 - torch.mean(loss)

        return loss


# class SoftIoULoss(nn.Module):
#     def __init__(self):
#         super(SoftIoULoss, self).__init__()
#     def forward(self, preds, gt_masks):
#         if isinstance(preds, list) or isinstance(preds, tuple):
#             loss_total = 0
#             for i in range(len(preds)):
#                 pred = preds[i]
#                 smooth = 1
#                 intersection = pred * gt_masks
#                 loss = (intersection.sum() + smooth) / (pred.sum() + gt_masks.sum() -intersection.sum() + smooth)
#                 loss = 1 - loss.mean()
#                 loss_total = loss_total + loss
#             return loss_total / len(preds)
#         else:
#             pred = preds
#             smooth = 1
#             intersection = pred * gt_masks
#             loss = (intersection.sum() + smooth) / (pred.sum() + gt_masks.sum() -intersection.sum() + smooth)
#             loss = 1 - loss.mean()
#             return loss

# class ISNetLoss(nn.Module):
#     def __init__(self):
#         super(ISNetLoss, self).__init__()
#         self.softiou = SoftIoULoss()
#         self.bce = nn.BCELoss()
#         self.grad = Get_gradient_nopadding()
#
#     def forward(self, preds, gt_masks):
#         edge_gt = self.grad(gt_masks.clone())
#
#         ### img loss
#         loss_img = self.softiou(preds[0], gt_masks)
#
#         ### edge loss
#         loss_edge = 10 * self.bce(preds[1], edge_gt)+ self.softiou(preds[1].sigmoid(), edge_gt)
#
#         return loss_img + loss_edge
