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.triton_helpers import libdevice
import torch.nn as nn
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def triton_per_fused_add_div_mul_pow_rsub_sum_0(in_out_ptr1, in_ptr0,
    in_ptr1, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 64
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex % 16
    r1 = rindex // 16
    tmp0 = tl.load(in_ptr0 + (32 + r0 + 64 * r1), None)
    tmp1 = tl.load(in_ptr1 + (32 + r0 + 64 * r1), None)
    tmp17 = tl.load(in_ptr0 + (r0 + 64 * r1), None)
    tmp18 = tl.load(in_ptr1 + (r0 + 64 * r1), None)
    tmp33 = tl.load(in_ptr0 + (48 + r0 + 64 * r1), None)
    tmp34 = tl.load(in_ptr1 + (48 + r0 + 64 * r1), None)
    tmp49 = tl.load(in_ptr0 + (16 + r0 + 64 * r1), None)
    tmp50 = tl.load(in_ptr1 + (16 + r0 + 64 * r1), None)
    tmp2 = tmp0 * tmp1
    tmp3 = tl.broadcast_to(tmp2, [XBLOCK, RBLOCK])
    tmp5 = tl.sum(tmp3, 1)[:, None]
    tmp6 = 1.0
    tmp7 = tmp6 - tmp0
    tmp8 = tmp7 * tmp1
    tmp9 = tl.broadcast_to(tmp8, [XBLOCK, RBLOCK])
    tmp11 = tl.sum(tmp9, 1)[:, None]
    tmp12 = tmp6 - tmp1
    tmp13 = tmp0 * tmp12
    tmp14 = tl.broadcast_to(tmp13, [XBLOCK, RBLOCK])
    tmp16 = tl.sum(tmp14, 1)[:, None]
    tmp19 = tmp17 * tmp18
    tmp20 = tl.broadcast_to(tmp19, [XBLOCK, RBLOCK])
    tmp22 = tl.sum(tmp20, 1)[:, None]
    tmp23 = tmp6 - tmp17
    tmp24 = tmp23 * tmp18
    tmp25 = tl.broadcast_to(tmp24, [XBLOCK, RBLOCK])
    tmp27 = tl.sum(tmp25, 1)[:, None]
    tmp28 = tmp6 - tmp18
    tmp29 = tmp17 * tmp28
    tmp30 = tl.broadcast_to(tmp29, [XBLOCK, RBLOCK])
    tmp32 = tl.sum(tmp30, 1)[:, None]
    tmp35 = tmp33 * tmp34
    tmp36 = tl.broadcast_to(tmp35, [XBLOCK, RBLOCK])
    tmp38 = tl.sum(tmp36, 1)[:, None]
    tmp39 = tmp6 - tmp33
    tmp40 = tmp39 * tmp34
    tmp41 = tl.broadcast_to(tmp40, [XBLOCK, RBLOCK])
    tmp43 = tl.sum(tmp41, 1)[:, None]
    tmp44 = tmp6 - tmp34
    tmp45 = tmp33 * tmp44
    tmp46 = tl.broadcast_to(tmp45, [XBLOCK, RBLOCK])
    tmp48 = tl.sum(tmp46, 1)[:, None]
    tmp51 = tmp49 * tmp50
    tmp52 = tl.broadcast_to(tmp51, [XBLOCK, RBLOCK])
    tmp54 = tl.sum(tmp52, 1)[:, None]
    tmp55 = tmp6 - tmp49
    tmp56 = tmp55 * tmp50
    tmp57 = tl.broadcast_to(tmp56, [XBLOCK, RBLOCK])
    tmp59 = tl.sum(tmp57, 1)[:, None]
    tmp60 = tmp6 - tmp50
    tmp61 = tmp49 * tmp60
    tmp62 = tl.broadcast_to(tmp61, [XBLOCK, RBLOCK])
    tmp64 = tl.sum(tmp62, 1)[:, None]
    tmp65 = tmp22 + tmp6
    tmp66 = 0.7
    tmp67 = tmp27 * tmp66
    tmp68 = tmp22 + tmp67
    tmp69 = 0.3
    tmp70 = tmp32 * tmp69
    tmp71 = tmp68 + tmp70
    tmp72 = tmp71 + tmp6
    tmp73 = tmp65 / tmp72
    tmp74 = tmp6 - tmp73
    tmp75 = 0.7518796992481203
    tmp76 = libdevice.pow(tmp74, tmp75)
    tmp77 = 0.0
    tmp78 = tmp76 + tmp77
    tmp79 = tmp54 + tmp6
    tmp80 = tmp59 * tmp66
    tmp81 = tmp54 + tmp80
    tmp82 = tmp64 * tmp69
    tmp83 = tmp81 + tmp82
    tmp84 = tmp83 + tmp6
    tmp85 = tmp79 / tmp84
    tmp86 = tmp6 - tmp85
    tmp87 = libdevice.pow(tmp86, tmp75)
    tmp88 = tmp78 + tmp87
    tmp89 = tmp5 + tmp6
    tmp90 = tmp11 * tmp66
    tmp91 = tmp5 + tmp90
    tmp92 = tmp16 * tmp69
    tmp93 = tmp91 + tmp92
    tmp94 = tmp93 + tmp6
    tmp95 = tmp89 / tmp94
    tmp96 = tmp6 - tmp95
    tmp97 = libdevice.pow(tmp96, tmp75)
    tmp98 = tmp88 + tmp97
    tmp99 = tmp38 + tmp6
    tmp100 = tmp43 * tmp66
    tmp101 = tmp38 + tmp100
    tmp102 = tmp48 * tmp69
    tmp103 = tmp101 + tmp102
    tmp104 = tmp103 + tmp6
    tmp105 = tmp99 / tmp104
    tmp106 = tmp6 - tmp105
    tmp107 = libdevice.pow(tmp106, tmp75)
    tmp108 = tmp98 + tmp107
    tmp109 = 0.25
    tmp110 = tmp108 * tmp109
    tl.debug_barrier()
    tl.store(in_out_ptr1 + tl.full([XBLOCK, 1], 0, tl.int32), tmp110, None)


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)
        buf10 = empty_strided_cuda((), (), torch.float32)
        buf13 = buf10
        del buf10
        buf14 = buf13
        del buf13
        get_raw_stream(0)
        triton_per_fused_add_div_mul_pow_rsub_sum_0[grid(1)](buf14, arg1_1,
            arg0_1, 1, 64, XBLOCK=1, num_warps=2, num_stages=1)
        del arg0_1
        del arg1_1
    return buf14,


class TverskyLoss(nn.Module):
    """Tversky Loss.

    .. seealso::
        Salehi, Seyed Sadegh Mohseni, Deniz Erdogmus, and Ali Gholipour. "Tversky loss function for image segmentation
        using 3D fully convolutional deep networks." International Workshop on Machine Learning in Medical Imaging.
        Springer, Cham, 2017.

    Args:
        alpha (float): Weight of false positive voxels.
        beta  (float): Weight of false negative voxels.
        smooth (float): Epsilon to avoid division by zero, when both Numerator and Denominator of Tversky are zeros.

    Attributes:
        alpha (float): Weight of false positive voxels.
        beta  (float): Weight of false negative voxels.
        smooth (float): Epsilon to avoid division by zero, when both Numerator and Denominator of Tversky are zeros.

    Notes:
        - setting alpha=beta=0.5: Equivalent to DiceLoss.
        - default parameters were suggested by https://arxiv.org/pdf/1706.05721.pdf .
    """

    def __init__(self, alpha=0.7, beta=0.3, smooth=1.0):
        super(TverskyLoss, self).__init__()
        self.alpha = alpha
        self.beta = beta
        self.smooth = smooth

    def tversky_index(self, y_pred, y_true):
        """Compute Tversky index.

        Args:
            y_pred (torch Tensor): Prediction.
            y_true (torch Tensor): Target.

        Returns:
            float: Tversky index.
        """
        y_true = y_true.float()
        tp = torch.sum(y_true * y_pred)
        fn = torch.sum(y_true * (1 - y_pred))
        fp = torch.sum((1 - y_true) * y_pred)
        numerator = tp + self.smooth
        denominator = tp + self.alpha * fp + self.beta * fn + self.smooth
        tversky_label = numerator / denominator
        return tversky_label

    def forward(self, input, target):
        n_classes = input.shape[1]
        tversky_sum = 0.0
        for i_label in range(n_classes):
            y_pred, y_true = input[:, i_label], target[:, i_label]
            tversky_sum += self.tversky_index(y_pred, y_true)
        return -tversky_sum / n_classes


class FocalTverskyLossNew(TverskyLoss):
    """Focal Tversky Loss.

    .. seealso::
        Abraham, Nabila, and Naimul Mefraz Khan. "A novel focal tversky loss function with improved attention u-net for
        lesion segmentation." 2019 IEEE 16th International Symposium on Biomedical Imaging (ISBI 2019). IEEE, 2019.

    Args:
        alpha (float): Weight of false positive voxels.
        beta  (float): Weight of false negative voxels.
        gamma (float): Typically between 1 and 3. Control between easy background and hard ROI training examples.
        smooth (float): Epsilon to avoid division by zero, when both Numerator and Denominator of Tversky are zeros.

    Attributes:
        gamma (float): Typically between 1 and 3. Control between easy background and hard ROI training examples.

    Notes:
        - setting alpha=beta=0.5 and gamma=1: Equivalent to DiceLoss.
        - default parameters were suggested by https://arxiv.org/pdf/1810.07842.pdf .
    """

    def __init__(self, alpha=0.7, beta=0.3, gamma=1.33, smooth=1.0):
        super(FocalTverskyLossNew, self).__init__()
        self.gamma = gamma
        self.tversky = TverskyLoss(alpha=alpha, beta=beta, smooth=smooth)

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


Measured latency_ms: 0.015479