import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, pred_hm, gt_hm, pred_wh, gt_wh, pred_reg, gt_reg, mask):
        pred_hm = torch.clamp(pred_hm, 1e-6, 1 - 1e-6)

        pos_inds = gt_hm.eq(1).float()
        neg_inds = gt_hm.lt(1).float()
        neg_weights = torch.pow(1 - gt_hm, 4)

        pos_loss = torch.log(pred_hm) * torch.pow(1 - pred_hm, 2) * pos_inds
        neg_loss = torch.log(1 - pred_hm) * torch.pow(pred_hm, 2) * neg_weights * neg_inds

        num_pos = pos_inds.sum()
        pos_loss_sum = pos_loss.sum()
        neg_loss_sum = neg_loss.sum()

        if num_pos > 0:
            hm_loss = -(pos_loss_sum + neg_loss_sum) / num_pos
        else:
            hm_loss = -neg_loss_sum

        mask_expanded = mask.expand_as(pred_wh)
        wh_loss = torch.sum(torch.abs(pred_wh - gt_wh) * mask_expanded)
        reg_loss = torch.sum(torch.abs(pred_reg - gt_reg) * mask_expanded)

        if num_pos > 0:
            wh_loss = wh_loss / num_pos
            reg_loss = reg_loss / num_pos

        return hm_loss + 0.1 * wh_loss + 1.0 * reg_loss


batch_size = 4
channels = 4
height = 128
width = 128


def get_inputs():
    pred_hm = torch.sigmoid(torch.randn(batch_size, channels, height, width))
    gt_hm = torch.bernoulli(torch.full((batch_size, channels, height, width), 0.1))
    pred_wh = torch.randn(batch_size, 2, height, width)
    gt_wh = torch.randn(batch_size, 2, height, width)
    pred_reg = torch.randn(batch_size, 2, height, width)
    gt_reg = torch.randn(batch_size, 2, height, width)
    mask = torch.bernoulli(torch.full((batch_size, 1, height, width), 0.1))
    return [pred_hm, gt_hm, pred_wh, gt_wh, pred_reg, gt_reg, mask]


def get_init_inputs():
    return []