'''
 * @author Waldinsamkeit
 * @email Zenglz_pro@163.com
 * @create date 2020-09-25 14:33:38
 * @desc 
'''

import torch


'''--------------------- Weighted Binary cross Entropy ----------------------'''

'''
In Torch BCELoss, weight is set to every element of input instead of to every class
'''
def weighted_binary_cross_entropy(output, target, weights=None):
    if weights is not None:
        assert len(weights) == 2
        
        loss = weights[1] * (target * torch.log(output)) + \
            weights[0] * ((1 - target) * torch.log(1 - output))
    else:
        loss = target * torch.log(output) + (1 - target) * torch.log(1 - output)

    return torch.neg(torch.mean(loss))


''' ---------------------- Binary focal loss function -------------------------- '''

'''
In some degree, it can reduce the influence of imbalanced dataset
'''

def focal_loss(y_true,y_pred,device):
    alpha,gamma = torch.tensor(0.25).to(device) , torch.tensor(2.0).to(device)
    y_pred=torch.clamp(y_pred,1e-7,1-1e-7)
    return - alpha * y_true * torch.log(y_pred) * (1 - y_pred) ** gamma\
        - (1 - alpha) * (1 - y_true) * torch.log(1 -  y_pred) * y_pred


[(36, 'torch.clamp', 'torch.clamp', (['y_pred', '(1e-07)', '(1 - 1e-07)'], {}), False, 'import torch\n'), (25, 'torch.mean', 'torch.mean', (['loss'], {}), False, 'import torch\n'), (23, 'torch.log', 'torch.log', (['output'], {}), False, 'import torch\n'), (23, 'torch.log', 'torch.log', (['(1 - output)'], {}), False, 'import torch\n'), (35, 'torch.tensor', 'torch.tensor', (['(0.25)'], {}), False, 'import torch\n'), (35, 'torch.tensor', 'torch.tensor', (['(2.0)'], {}), False, 'import torch\n'), (37, 'torch.log', 'torch.log', (['y_pred'], {}), False, 'import torch\n'), (38, 'torch.log', 'torch.log', (['(1 - y_pred)'], {}), False, 'import torch\n'), (20, 'torch.log', 'torch.log', (['output'], {}), False, 'import torch\n'), (21, 'torch.log', 'torch.log', (['(1 - output)'], {}), False, 'import torch\n')]