import torch


def weighted_binary_cross_entropy2(sigmoid_x, y, weighted_matrix, weight=None, reduction=None):
    """
    Aha this is correct!
    sigmoid_x = nn.Sigmoid()(x)
    Args:
        sigmoid_x: predicted probability of size [N,C], N sample and C Class. Eg. Must be in range of [0,1]
        targets: true value, one-hot-like vector of size [N,C]
        pos_weight: Weight for postive sample
    """
    if not (y.size() == sigmoid_x.size()):
        raise ValueError("Target size ({}) must be the same as input size ({})".format(y.size(), sigmoid_x.size()))
   
    #print("y.size(), sigmoid_x.size()", y.size(), sigmoid_x.size())
    sigmoid_x = torch.clamp(sigmoid_x,min=1e-7,max=1-1e-7) 
    loss = - torch.matmul(y*sigmoid_x.log() + (1-y)*(1-sigmoid_x).log(), weighted_matrix)
    
    if weight is not None:
        loss = loss * weight
        
    if reduction is None:
        return loss
    elif reduction == 'mean':
        return loss.mean()
    elif reduction == 'sum':
        return loss.sum()
    return None