import torch


def focal_l2_loss(s, sxing, gamma=1, multi_task_weight=0.1,
                  keypoint_task_weight=1, nstack_weight=[1, 1, 1, 1], alpha=0., beta=0.):
    """
    Compute the focal L2 loss between predicted and groundtruth score maps.
    :param s:  predicted tensor (nstack, batch, channel, height, width), predicted score maps
    :param sxing: target tensor (nstack, batch, channel, height, width)
    :param mask_miss: tensor (nstack, batch, 1, height, width)
    :param gamma: focusing parameter
    :return: a scalar tensor
    """
    # eps = 1e-8  # 1e-12
    # s = torch.clamp(s, eps, 1. - eps)  # improve the stability of the focal loss
    # mask = mask_miss.expand_as(sxing).clone()  # type: torch.Tensor
    # del mask_miss
    # mask[:, :, -2, :, :] *= multi_task_weight  # except for person mask channel
    # mask[:, :, heat_start:bkg_start, :, :] *= keypoint_task_weight
    # pdb.set_trace()
    st = torch.where(torch.ge(sxing, 0.01), s - alpha, 1 - s - beta)  
    factor = torch.abs(1. - st)  # (1. - st) ** gamma  for gamma=2
    # multiplied by mask_miss via broadcast operation
    out = (s - sxing) ** 2 * factor # * mask  # type: torch.Tensor
    # sum over the feature map, should divide by batch afterwards
    # loss_nstack = out.sum(dim=(1, 2, 3, 4))   # out.[:, :, heat_start:bkg_start, :, :] 
    loss_nstack = out.mean(dim=(1, 2, 3, 4))   # out.[:, :, heat_start:bkg_start, :, :] 
    assert len(loss_nstack) == len(nstack_weight), nstack_weight
    # print(' heatmap focal L2 loss per stack..........  ', loss_nstack.detach().cpu().numpy())
    weight_loss = [loss_nstack[i] * nstack_weight[i] for i in range(len(nstack_weight))]
    loss = sum(weight_loss) / sum(nstack_weight)  
    # pdb.set_trace()
    return loss