'''
 * @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, 11, 36, 42), 'torch.clamp', 'torch.clamp', ({(36, 23, 36, 29): 'y_pred', (36, 30, 36, 34): '1e-07', (36, 35, 36, 41): '1 - 1e-07'}, {}), '(y_pred, 1e-07, 1 - 1e-07)', False, 'import torch\n'), ((25, 21, 25, 37), 'torch.mean', 'torch.mean', ({(25, 32, 25, 36): 'loss'}, {}), '(loss)', False, 'import torch\n'), ((23, 24, 23, 41), 'torch.log', 'torch.log', ({(23, 34, 23, 40): 'output'}, {}), '(output)', False, 'import torch\n'), ((23, 59, 23, 80), 'torch.log', 'torch.log', ({(23, 69, 23, 79): '(1 - output)'}, {}), '(1 - output)', False, 'import torch\n'), ((35, 18, 35, 36), 'torch.tensor', 'torch.tensor', ({(35, 31, 35, 35): '(0.25)'}, {}), '(0.25)', False, 'import torch\n'), ((35, 50, 35, 67), 'torch.tensor', 'torch.tensor', ({(35, 63, 35, 66): '(2.0)'}, {}), '(2.0)', False, 'import torch\n'), ((37, 30, 37, 47), 'torch.log', 'torch.log', ({(37, 40, 37, 46): 'y_pred'}, {}), '(y_pred)', False, 'import torch\n'), ((38, 39, 38, 61), 'torch.log', 'torch.log', ({(38, 49, 38, 60): '(1 - y_pred)'}, {}), '(1 - y_pred)', False, 'import torch\n'), ((20, 38, 20, 55), 'torch.log', 'torch.log', ({(20, 48, 20, 54): 'output'}, {}), '(output)', False, 'import torch\n'), ((21, 41, 21, 62), 'torch.log', 'torch.log', ({(21, 51, 21, 61): '(1 - output)'}, {}), '(1 - output)', False, 'import torch\n')]