
# pairwise loss functions in AUROC benchmark work



import torch 
import torch.nn.functional as F

 

    
class PSQLoss(torch.nn.Module): # pairwise squared loss
    def __init__(self, margin=1.0): 
        super(PSQLoss, self).__init__() 
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.margin = margin

    def forward(self, y_pred, y_true): 
        f_ps = y_pred[y_true == 1].reshape(-1, 1)
        f_ns = y_pred[y_true == 0].reshape(-1, 1)
        if len(f_ns) == 0 or len(f_ps) == 0:
            print('switching to logitstic loss')
            criterion = F.binary_cross_entropy_with_logits  # with sigmoid
            return criterion(y_pred, y_true)
        f_ps = f_ps.repeat(1,len(f_ns))
        f_ns = f_ns.repeat(1,len(f_ps))
        difference = f_ps - f_ns.transpose(0,1)
        difference = difference - self.margin
        difference = difference ** 2
        #loss = difference.mean()  
        loss = difference
        return loss

    
class PSMLoss(torch.nn.Module): # pairwise sigmoid loss
    def __init__(self, margin=1.0): 
        super(PSMLoss, self).__init__() 
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.margin = margin

    def forward(self, y_pred, y_true): 
        y_pred = torch.clamp(y_pred, min=-3, max=3)  
        f_ps = y_pred[y_true == 1].reshape(-1, 1)
        f_ns = y_pred[y_true == 0].reshape(-1, 1)
        if len(f_ns) == 0 or len(f_ps) == 0:
            print('switching to logitstic loss')
            criterion = F.binary_cross_entropy_with_logits  # with sigmoid
            return criterion(y_pred, y_true)
        f_ps = f_ps.repeat(1,len(f_ns))
        f_ns = f_ns.repeat(1,len(f_ps))
        difference = f_ps - f_ns.transpose(0,1)
        difference = 1 + torch.exp(self.margin*difference)
        difference = 1./difference
        loss = difference.mean()  
        return loss

    
class PHLoss(torch.nn.Module): # pairwise hinge loss
    def __init__(self, margin=1.0): 
        super(PHLoss, self).__init__() 
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.margin = margin

    def forward(self, y_pred, y_true): 
        f_ps = y_pred[y_true == 1].reshape(-1, 1)
        f_ns = y_pred[y_true == 0].reshape(-1, 1)
        if len(f_ns) == 0 or len(f_ps) == 0:
            print('switching to logitstic loss')
            criterion = F.binary_cross_entropy_with_logits  # with sigmoid
            return criterion(y_pred, y_true)
        f_ps = f_ps.repeat(1,len(f_ns))
        f_ns = f_ns.repeat(1,len(f_ps))
        difference = f_ps - f_ns.transpose(0,1)
        difference = torch.maximum(self.margin - difference, torch.tensor(0.0))
        loss = difference.mean()  
        return loss

    
class PBHLoss(torch.nn.Module): # pairwise barrier hinge loss
    def __init__(self, r=1.0, b=2.0): 
        super(PBHLoss, self).__init__() 
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.r = r
        self.b = b

    def forward(self, y_pred, y_true): 
        f_ps = y_pred[y_true == 1].reshape(-1, 1)
        f_ns = y_pred[y_true == 0].reshape(-1, 1)
        if len(f_ns) == 0 or len(f_ps) == 0:
            print('switching to logitstic loss')
            criterion = F.binary_cross_entropy_with_logits  # with sigmoid
            return criterion(y_pred, y_true)
        f_ps = f_ps.repeat(1,len(f_ns))
        f_ns = f_ns.repeat(1,len(f_ps))
        difference = f_ps - f_ns.transpose(0,1)
        loss = torch.maximum(self.r - difference, self.b * (difference - self.r))
        loss = torch.maximum(-self.b * (self.r + difference) + self.r, loss)
        loss = loss.mean()  
        return loss

    
class PSHLoss(torch.nn.Module): # pairwise squared hinge loss
    def __init__(self, margin=1.0): 
        super(PSHLoss, self).__init__() 
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.margin = margin

    def forward(self, y_pred, y_true): 
        f_ps = y_pred[y_true == 1].view(-1, 1)
        f_ns = y_pred[y_true == 0].view(-1, 1)
        if len(f_ns) == 0 or len(f_ps) == 0:
            print('switching to logitstic loss')
            criterion = F.binary_cross_entropy_with_logits  # with sigmoid
            return criterion(y_pred, y_true)
        f_ps = f_ps.repeat(1,len(f_ns))
        f_ns = f_ns.repeat(1,len(f_ps))
        difference = f_ps - f_ns.transpose(0,1)
        difference = torch.maximum(self.margin - difference, torch.tensor(0.0))
        difference = difference ** 2
        loss = difference.mean()  
        return loss

 
 
class PLLoss(torch.nn.Module): # pairwise logistic loss.
    def __init__(self, margin=1.0): 
        super(PLLoss, self).__init__() 
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.margin = margin

    def forward(self, y_pred, y_true): 
        f_ps = y_pred[y_true == 1].reshape(-1, 1)
        f_ns = y_pred[y_true == 0].reshape(-1, 1)
        if len(f_ns) == 0 or len(f_ps) == 0:
            print('switching to logitstic loss')
            criterion = F.binary_cross_entropy_with_logits  # with sigmoid
            return criterion(y_pred, y_true)
        f_ps = f_ps.repeat(1,len(f_ns))
        f_ns = f_ns.repeat(1,len(f_ps))
        difference = f_ps - f_ns.transpose(0,1)
        difference = 1 + torch.exp(-self.margin*difference)
        difference = torch.log(difference)
        loss = difference.mean()  
        return loss




