import torch
import torch.distributed
import torch
import torch.nn as nn
import torch.nn.functional
import torch.utils.data
import torch.optim
import torch.optim.lr_scheduler


def bounded_iou_loss(pred, target, beta=0.2, eps=0.001):
    """BIoULoss.

    This is an implementation of paper
    `Improving Object Localization with Fitness NMS and Bounded IoU Loss.
    <https://arxiv.org/abs/1711.00164>`_.

    Args:
        pred (torch.Tensor): Predicted bboxes.
        target (torch.Tensor): Target bboxes.
        beta (float): beta parameter in smoothl1.
        eps (float): eps to avoid NaN.
    """
    pred_ctrx = (pred[:, 0] + pred[:, 2]) * 0.5
    pred_ctry = (pred[:, 1] + pred[:, 3]) * 0.5
    pred_w = pred[:, 2] - pred[:, 0]
    pred_h = pred[:, 3] - pred[:, 1]
    with torch.no_grad():
        target_ctrx = (target[:, 0] + target[:, 2]) * 0.5
        target_ctry = (target[:, 1] + target[:, 3]) * 0.5
        target_w = target[:, 2] - target[:, 0]
        target_h = target[:, 3] - target[:, 1]
    dx = target_ctrx - pred_ctrx
    dy = target_ctry - pred_ctry
    loss_dx = 1 - torch.max((target_w - 2 * dx.abs()) / (target_w + 2 * dx.
        abs() + eps), torch.zeros_like(dx))
    loss_dy = 1 - torch.max((target_h - 2 * dy.abs()) / (target_h + 2 * dy.
        abs() + eps), torch.zeros_like(dy))
    loss_dw = 1 - torch.min(target_w / (pred_w + eps), pred_w / (target_w +
        eps))
    loss_dh = 1 - torch.min(target_h / (pred_h + eps), pred_h / (target_h +
        eps))
    loss_comb = torch.stack([loss_dx, loss_dy, loss_dw, loss_dh], dim=-1).view(
        loss_dx.size(0), -1)
    loss = torch.where(loss_comb < beta, 0.5 * loss_comb * loss_comb / beta,
        loss_comb - 0.5 * beta)
    return loss


class BoundedIoULoss(nn.Module):

    def __init__(self, beta=0.2, eps=0.001):
        super(BoundedIoULoss, self).__init__()
        self.beta = beta
        self.eps = eps

    def forward(self, pred, target):
        return bounded_iou_loss(pred, target, self.beta, self.eps)


def get_inputs():
    return [torch.rand([4, 4, 4, 4]), torch.rand([4, 4, 4, 4])]


def get_init_inputs():
    return [[], {}]

import torch
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import math as tl_math
import torch.distributed
import torch
import torch.nn as nn
import torch.nn.functional
import torch.utils.data
import torch.optim
import torch.optim.lr_scheduler
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor


@triton.jit
def triton_poi_fused_div_lt_mul_stack_sub_where_0(in_out_ptr0, in_ptr0,
    in_ptr1, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4 % 16
    x2 = xindex // 64
    x3 = xindex
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (32 + x1 + 64 * x2), tmp4 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp6 = tl.load(in_ptr0 + (x1 + 64 * x2), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp7 = tmp5 - tmp6
    tmp8 = tmp6 + tmp5
    tmp9 = 0.5
    tmp10 = tmp8 * tmp9
    tmp11 = tl.load(in_ptr1 + (x1 + 64 * x2), tmp4 & xmask, eviction_policy
        ='evict_last', other=0.0)
    tmp12 = tl.load(in_ptr1 + (32 + x1 + 64 * x2), tmp4 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp13 = tmp11 + tmp12
    tmp14 = tmp13 * tmp9
    tmp15 = tmp10 - tmp14
    tmp16 = tl_math.abs(tmp15)
    tmp17 = 2.0
    tmp18 = tmp16 * tmp17
    tmp19 = tmp7 - tmp18
    tmp20 = tmp7 + tmp18
    tmp21 = 0.001
    tmp22 = tmp20 + tmp21
    tmp23 = tmp19 / tmp22
    tmp24 = 0.0
    tmp25 = triton_helpers.maximum(tmp23, tmp24)
    tmp26 = 1.0
    tmp27 = tmp26 - tmp25
    tmp28 = tl.full(tmp27.shape, 0.0, tmp27.dtype)
    tmp29 = tl.where(tmp4, tmp27, tmp28)
    tmp30 = tmp0 >= tmp3
    tmp31 = tl.full([1], 2, tl.int64)
    tmp32 = tmp0 < tmp31
    tmp33 = tmp30 & tmp32
    tmp34 = tl.load(in_ptr0 + (48 + x1 + 64 * x2), tmp33 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp35 = tl.load(in_ptr0 + (16 + x1 + 64 * x2), tmp33 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp36 = tmp34 - tmp35
    tmp37 = tmp35 + tmp34
    tmp38 = tmp37 * tmp9
    tmp39 = tl.load(in_ptr1 + (16 + x1 + 64 * x2), tmp33 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp40 = tl.load(in_ptr1 + (48 + x1 + 64 * x2), tmp33 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp41 = tmp39 + tmp40
    tmp42 = tmp41 * tmp9
    tmp43 = tmp38 - tmp42
    tmp44 = tl_math.abs(tmp43)
    tmp45 = tmp44 * tmp17
    tmp46 = tmp36 - tmp45
    tmp47 = tmp36 + tmp45
    tmp48 = tmp47 + tmp21
    tmp49 = tmp46 / tmp48
    tmp50 = triton_helpers.maximum(tmp49, tmp24)
    tmp51 = tmp26 - tmp50
    tmp52 = tl.full(tmp51.shape, 0.0, tmp51.dtype)
    tmp53 = tl.where(tmp33, tmp51, tmp52)
    tmp54 = tmp0 >= tmp31
    tmp55 = tl.full([1], 3, tl.int64)
    tmp56 = tmp0 < tmp55
    tmp57 = tmp54 & tmp56
    tmp58 = tl.load(in_ptr0 + (32 + x1 + 64 * x2), tmp57 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp59 = tl.load(in_ptr0 + (x1 + 64 * x2), tmp57 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp60 = tmp58 - tmp59
    tmp61 = tl.load(in_ptr1 + (32 + x1 + 64 * x2), tmp57 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp62 = tl.load(in_ptr1 + (x1 + 64 * x2), tmp57 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp63 = tmp61 - tmp62
    tmp64 = tmp63 + tmp21
    tmp65 = tmp60 / tmp64
    tmp66 = tmp60 + tmp21
    tmp67 = tmp63 / tmp66
    tmp68 = triton_helpers.minimum(tmp65, tmp67)
    tmp69 = tmp26 - tmp68
    tmp70 = tl.full(tmp69.shape, 0.0, tmp69.dtype)
    tmp71 = tl.where(tmp57, tmp69, tmp70)
    tmp72 = tmp0 >= tmp55
    tl.full([1], 4, tl.int64)
    tmp75 = tl.load(in_ptr0 + (48 + x1 + 64 * x2), tmp72 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp76 = tl.load(in_ptr0 + (16 + x1 + 64 * x2), tmp72 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp77 = tmp75 - tmp76
    tmp78 = tl.load(in_ptr1 + (48 + x1 + 64 * x2), tmp72 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp79 = tl.load(in_ptr1 + (16 + x1 + 64 * x2), tmp72 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp80 = tmp78 - tmp79
    tmp81 = tmp80 + tmp21
    tmp82 = tmp77 / tmp81
    tmp83 = tmp77 + tmp21
    tmp84 = tmp80 / tmp83
    tmp85 = triton_helpers.minimum(tmp82, tmp84)
    tmp86 = tmp26 - tmp85
    tmp87 = tl.full(tmp86.shape, 0.0, tmp86.dtype)
    tmp88 = tl.where(tmp72, tmp86, tmp87)
    tmp89 = tl.where(tmp57, tmp71, tmp88)
    tmp90 = tl.where(tmp33, tmp53, tmp89)
    tmp91 = tl.where(tmp4, tmp29, tmp90)
    tmp92 = 0.2
    tmp93 = tmp91 < tmp92
    tmp94 = tmp91 * tmp9
    tmp95 = tmp94 * tmp91
    tmp96 = 5.0
    tmp97 = tmp95 * tmp96
    tmp98 = 0.1
    tmp99 = tmp91 - tmp98
    tmp100 = tl.where(tmp93, tmp97, tmp99)
    tl.store(in_out_ptr0 + x3, tmp100, xmask)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(arg1_1, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        buf1 = reinterpret_tensor(buf0, (4, 64), (64, 1), 0)
        del buf0
        get_raw_stream(0)
        triton_poi_fused_div_lt_mul_stack_sub_where_0[grid(256)](buf1,
            arg1_1, arg0_1, 256, XBLOCK=128, num_warps=4, num_stages=1)
        del arg0_1
        del arg1_1
    return buf1,


def bounded_iou_loss(pred, target, beta=0.2, eps=0.001):
    """BIoULoss.

    This is an implementation of paper
    `Improving Object Localization with Fitness NMS and Bounded IoU Loss.
    <https://arxiv.org/abs/1711.00164>`_.

    Args:
        pred (torch.Tensor): Predicted bboxes.
        target (torch.Tensor): Target bboxes.
        beta (float): beta parameter in smoothl1.
        eps (float): eps to avoid NaN.
    """
    pred_ctrx = (pred[:, 0] + pred[:, 2]) * 0.5
    pred_ctry = (pred[:, 1] + pred[:, 3]) * 0.5
    pred_w = pred[:, 2] - pred[:, 0]
    pred_h = pred[:, 3] - pred[:, 1]
    with torch.no_grad():
        target_ctrx = (target[:, 0] + target[:, 2]) * 0.5
        target_ctry = (target[:, 1] + target[:, 3]) * 0.5
        target_w = target[:, 2] - target[:, 0]
        target_h = target[:, 3] - target[:, 1]
    dx = target_ctrx - pred_ctrx
    dy = target_ctry - pred_ctry
    loss_dx = 1 - torch.max((target_w - 2 * dx.abs()) / (target_w + 2 * dx.
        abs() + eps), torch.zeros_like(dx))
    loss_dy = 1 - torch.max((target_h - 2 * dy.abs()) / (target_h + 2 * dy.
        abs() + eps), torch.zeros_like(dy))
    loss_dw = 1 - torch.min(target_w / (pred_w + eps), pred_w / (target_w +
        eps))
    loss_dh = 1 - torch.min(target_h / (pred_h + eps), pred_h / (target_h +
        eps))
    loss_comb = torch.stack([loss_dx, loss_dy, loss_dw, loss_dh], dim=-1).view(
        loss_dx.size(0), -1)
    loss = torch.where(loss_comb < beta, 0.5 * loss_comb * loss_comb / beta,
        loss_comb - 0.5 * beta)
    return loss


class BoundedIoULossNew(nn.Module):

    def __init__(self, beta=0.2, eps=0.001):
        super(BoundedIoULossNew, self).__init__()
        self.beta = beta
        self.eps = eps

    def forward(self, input_0, input_1):
        arg0_1 = input_0
        arg1_1 = input_1
        output = call([arg0_1, arg1_1])
        return output[0]
