import torch


def neg_loss(pred, gt):
    """Modified focal loss. Exactly the same as CornerNet.
    Runs faster and costs a little bit more memory

    Args:
        pred (torch.Tensor): predicted center heatmaps,
            should have shapes [batch, c, h, w]
        gt (torch.Tensor): ground truth center heatmaps,
            should have shapes [batch, c, h, w]

    Returns:
        torch.Tensor with focal loss value.
    """
    pred = pred.unsqueeze(1).float()
    gt = gt.unsqueeze(1).float()

    positive_inds = gt.eq(1).float()
    negative_inds = gt.lt(1).float()
    negative_weights = torch.pow(1 - gt, 4)

    loss = 0

    positive_loss = torch.log(pred + 1e-12) * torch.pow(1 - pred, 3) * positive_inds
    negative_loss = (
        torch.log(1 - pred + 1e-12) * torch.pow(pred, 3) * negative_weights * negative_inds
    )

    num_pos = positive_inds.float().sum()
    positive_loss = positive_loss.sum()
    negative_loss = negative_loss.sum()

    if num_pos == 0:
        loss = loss - negative_loss
    else:
        loss = loss - (positive_loss + negative_loss) / num_pos

    return loss