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


@triton.jit
def _triton_helper_fn_add0(arg0_0, arg1_0):
    tmp0 = arg0_0 + arg1_0
    return tmp0


@triton.jit
def triton_per_fused_cumsum_index_mul_rsub_sort_sub_sum_0(in_ptr0, in_ptr1,
    out_ptr0, out_ptr2, out_ptr3, out_ptr4, out_ptr5, out_ptr6, out_ptr7,
    xnumel, rnumel):
    XBLOCK: tl.constexpr = 1
    RBLOCK: tl.constexpr = 256
    xoffset = tl.program_id(0) * XBLOCK
    tl.full([1], xoffset, tl.int32)
    tl.full([RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[:]
    tl.full([RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + r0, None)
    tmp1 = tl.load(in_ptr1 + r0, None)
    tmp2 = 2.0
    tmp3 = tmp1 * tmp2
    tmp4 = 1.0
    tmp5 = tmp3 - tmp4
    tmp6 = tmp0 * tmp5
    tmp7 = tmp4 - tmp6
    tmp8 = r0
    tmp9 = tmp8.to(tl.int16)
    tmp10 = tl.broadcast_to(tmp7, [RBLOCK])
    tmp11 = tl.broadcast_to(tmp9, [RBLOCK])
    tmp12, tmp13 = triton_helpers.sort_with_index(tmp10, tmp11, None, 0,
        stable=False, descending=True)
    tmp14 = tmp0 * tmp1
    tmp15 = tl.broadcast_to(tmp14, [RBLOCK])
    tmp17 = triton_helpers.promote_to_tensor(tl.sum(tmp15, 0))
    tmp18 = tl.broadcast_to(tmp0, [RBLOCK])
    tmp20 = triton_helpers.promote_to_tensor(tl.sum(tmp18, 0))
    tmp21 = tl.broadcast_to(tmp1, [RBLOCK])
    tmp23 = triton_helpers.promote_to_tensor(tl.sum(tmp21, 0))
    tmp24 = tmp13.to(tl.int64)
    tmp25 = tl.full([RBLOCK], 256, tl.int32)
    tmp26 = tmp24 + tmp25
    tmp27 = tmp24 < 0
    tmp28 = tl.where(tmp27, tmp26, tmp24)
    tl.device_assert((0 <= tmp28) & (tmp28 < 256),
        'index out of bounds: 0 <= tmp28 < 256')
    tmp30 = tl.load(in_ptr1 + tmp28, None, eviction_policy='evict_last')
    tmp31 = tl.broadcast_to(tmp30, [RBLOCK])
    tmp33 = triton_helpers.promote_to_tensor(tl.sum(tmp31, 0))
    tmp34 = tmp30.to(tl.float32)
    tmp35 = tl.broadcast_to(tmp34, [RBLOCK])
    tmp36, = tl.associative_scan((tmp35,), 0, _triton_helper_fn_add0)
    tmp37 = tmp4 - tmp30
    tmp38 = tmp37.to(tl.float32)
    tmp39 = tl.broadcast_to(tmp38, [RBLOCK])
    tmp40, = tl.associative_scan((tmp39,), 0, _triton_helper_fn_add0)
    tl.store(out_ptr0 + tl.broadcast_to(r0, [RBLOCK]), tmp12, None)
    tl.store(out_ptr6 + tl.broadcast_to(r0, [RBLOCK]), tmp36, None)
    tl.store(out_ptr7 + tl.broadcast_to(r0, [RBLOCK]), tmp40, None)
    tl.store(out_ptr2 + tl.full([1], 0, tl.int32), tmp17, None)
    tl.store(out_ptr3 + tl.full([1], 0, tl.int32), tmp20, None)
    tl.store(out_ptr4 + tl.full([1], 0, tl.int32), tmp23, None)
    tl.store(out_ptr5 + tl.full([1], 0, tl.int32), tmp33, None)


@triton.jit
def triton_per_fused_add_clamp_log_mean_mul_neg_pow_rsub_sum_1(in_ptr0,
    in_ptr1, out_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 + (r0 + 64 * r1), None)
    tmp1 = tl.load(in_ptr1 + (r0 + 64 * r1), None)
    tmp22 = tl.load(in_ptr0 + (16 + r0 + 64 * r1), None)
    tmp23 = tl.load(in_ptr1 + (16 + r0 + 64 * r1), None)
    tmp42 = tl.load(in_ptr0 + (32 + r0 + 64 * r1), None)
    tmp43 = tl.load(in_ptr1 + (32 + r0 + 64 * r1), None)
    tmp62 = tl.load(in_ptr0 + (48 + r0 + 64 * r1), None)
    tmp63 = tl.load(in_ptr1 + (48 + r0 + 64 * r1), None)
    tmp2 = tmp0 * tmp1
    tmp3 = 1.0
    tmp4 = tmp3 - tmp0
    tmp5 = tmp3 - tmp1
    tmp6 = tmp4 * tmp5
    tmp7 = tmp2 + tmp6
    tmp8 = tmp3 - tmp7
    tmp9 = tmp8 * tmp8
    tmp10 = 1e-12
    tmp11 = triton_helpers.maximum(tmp0, tmp10)
    tmp12 = tl_math.log(tmp11)
    tmp13 = tmp1 * tmp12
    tmp14 = triton_helpers.maximum(tmp4, tmp10)
    tmp15 = tl_math.log(tmp14)
    tmp16 = tmp5 * tmp15
    tmp17 = tmp13 + tmp16
    tmp18 = -tmp17
    tmp19 = tmp9 * tmp18
    tmp20 = 0.25
    tmp21 = tmp19 * tmp20
    tmp24 = tmp22 * tmp23
    tmp25 = tmp3 - tmp22
    tmp26 = tmp3 - tmp23
    tmp27 = tmp25 * tmp26
    tmp28 = tmp24 + tmp27
    tmp29 = tmp3 - tmp28
    tmp30 = tmp29 * tmp29
    tmp31 = triton_helpers.maximum(tmp22, tmp10)
    tmp32 = tl_math.log(tmp31)
    tmp33 = tmp23 * tmp32
    tmp34 = triton_helpers.maximum(tmp25, tmp10)
    tmp35 = tl_math.log(tmp34)
    tmp36 = tmp26 * tmp35
    tmp37 = tmp33 + tmp36
    tmp38 = -tmp37
    tmp39 = tmp30 * tmp38
    tmp40 = tmp39 * tmp20
    tmp41 = tmp21 + tmp40
    tmp44 = tmp42 * tmp43
    tmp45 = tmp3 - tmp42
    tmp46 = tmp3 - tmp43
    tmp47 = tmp45 * tmp46
    tmp48 = tmp44 + tmp47
    tmp49 = tmp3 - tmp48
    tmp50 = tmp49 * tmp49
    tmp51 = triton_helpers.maximum(tmp42, tmp10)
    tmp52 = tl_math.log(tmp51)
    tmp53 = tmp43 * tmp52
    tmp54 = triton_helpers.maximum(tmp45, tmp10)
    tmp55 = tl_math.log(tmp54)
    tmp56 = tmp46 * tmp55
    tmp57 = tmp53 + tmp56
    tmp58 = -tmp57
    tmp59 = tmp50 * tmp58
    tmp60 = tmp59 * tmp20
    tmp61 = tmp41 + tmp60
    tmp64 = tmp62 * tmp63
    tmp65 = tmp3 - tmp62
    tmp66 = tmp3 - tmp63
    tmp67 = tmp65 * tmp66
    tmp68 = tmp64 + tmp67
    tmp69 = tmp3 - tmp68
    tmp70 = tmp69 * tmp69
    tmp71 = triton_helpers.maximum(tmp62, tmp10)
    tmp72 = tl_math.log(tmp71)
    tmp73 = tmp63 * tmp72
    tmp74 = triton_helpers.maximum(tmp65, tmp10)
    tmp75 = tl_math.log(tmp74)
    tmp76 = tmp66 * tmp75
    tmp77 = tmp73 + tmp76
    tmp78 = -tmp77
    tmp79 = tmp70 * tmp78
    tmp80 = tmp79 * tmp20
    tmp81 = tmp61 + tmp80
    tmp82 = tl.broadcast_to(tmp81, [XBLOCK, RBLOCK])
    tmp84 = tl.sum(tmp82, 1)[:, None]
    tl.store(out_ptr1 + tl.full([XBLOCK, 1], 0, tl.int32), tmp84, None)


@triton.jit
def triton_per_fused_add_div_dot_lift_fresh_mean_mul_relu_rsub_sub_2(
    in_out_ptr0, in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4, in_ptr5,
    in_ptr6, xnumel, rnumel):
    XBLOCK: tl.constexpr = 1
    RBLOCK: tl.constexpr = 256
    xoffset = tl.program_id(0) * XBLOCK
    tl.full([1], xoffset, tl.int32)
    tl.full([RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[:]
    tl.full([RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + r0, None)
    tmp6 = tl.load(in_out_ptr0 + 0)
    tmp7 = tl.broadcast_to(tmp6, [RBLOCK])
    tmp24 = tl.load(in_ptr1 + r0, None)
    tmp26 = tl.load(in_ptr2 + r0, None)
    tmp35 = tl.load(in_ptr3 + 0)
    tmp36 = tl.broadcast_to(tmp35, [1])
    tmp40 = tl.load(in_ptr4 + 0)
    tmp41 = tl.broadcast_to(tmp40, [1])
    tmp42 = tl.load(in_ptr5 + 0)
    tmp43 = tl.broadcast_to(tmp42, [1])
    tmp49 = tl.load(in_ptr6 + 0)
    tmp50 = tl.broadcast_to(tmp49, [1])
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = triton_helpers.maximum(tmp1, tmp0)
    tmp3 = r0
    tmp4 = tl.full([1], 1, tl.int64)
    tmp5 = tmp3 >= tmp4
    tmp8 = tl.load(in_ptr1 + tl.broadcast_to(r0, [RBLOCK]), tmp5, other=0.0)
    tmp9 = tmp7 - tmp8
    tmp10 = tl.load(in_ptr2 + tl.broadcast_to(r0, [RBLOCK]), tmp5, other=0.0)
    tmp11 = tmp7 + tmp10
    tmp12 = tmp9 / tmp11
    tmp13 = 1.0
    tmp14 = tmp13 - tmp12
    tmp15 = tl.load(in_ptr1 + tl.broadcast_to(-1 + r0, [RBLOCK]), tmp5,
        other=0.0)
    tmp16 = tmp7 - tmp15
    tmp17 = tl.load(in_ptr2 + tl.broadcast_to(-1 + r0, [RBLOCK]), tmp5,
        other=0.0)
    tmp18 = tmp7 + tmp17
    tmp19 = tmp16 / tmp18
    tmp20 = tmp13 - tmp19
    tmp21 = tmp14 - tmp20
    tmp22 = tl.full(tmp21.shape, 0.0, tmp21.dtype)
    tmp23 = tl.where(tmp5, tmp21, tmp22)
    tmp25 = tmp7 - tmp24
    tmp27 = tmp7 + tmp26
    tmp28 = tmp25 / tmp27
    tmp29 = tmp13 - tmp28
    tmp30 = tl.where(tmp5, tmp23, tmp29)
    tmp31 = tmp2 * tmp30
    tmp32 = tl.broadcast_to(tmp31, [RBLOCK])
    tmp34 = triton_helpers.promote_to_tensor(tl.sum(tmp32, 0))
    tmp37 = 2.0
    tmp38 = tmp36 * tmp37
    tmp39 = tmp38 + tmp13
    tmp44 = tmp41 + tmp43
    tmp45 = tmp44 + tmp13
    tmp46 = tmp39 / tmp45
    tmp47 = tmp13 - tmp46
    tmp48 = tmp47 * tmp13
    tmp51 = 64.0
    tmp52 = tmp50 / tmp51
    tmp53 = 0.2
    tmp54 = tmp52 * tmp53
    tmp55 = tmp48 + tmp54
    tmp56 = 0.0
    tmp57 = tmp34 * tmp56
    tmp58 = tmp55 + tmp57
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([1], 0, tl.int32), tmp58, 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)
        buf0 = empty_strided_cuda((256,), (1,), torch.float32)
        buf5 = empty_strided_cuda((), (), torch.float32)
        buf6 = empty_strided_cuda((), (), torch.float32)
        buf7 = empty_strided_cuda((), (), torch.float32)
        buf2 = empty_strided_cuda((), (), torch.float32)
        buf3 = empty_strided_cuda((256,), (1,), torch.float32)
        buf4 = empty_strided_cuda((256,), (1,), torch.float32)
        get_raw_stream(0)
        triton_per_fused_cumsum_index_mul_rsub_sort_sub_sum_0[grid(1)](arg0_1,
            arg1_1, buf0, buf5, buf6, buf7, buf2, buf3, buf4, 1, 256,
            num_warps=2, num_stages=1)
        buf9 = empty_strided_cuda((), (), torch.float32)
        triton_per_fused_add_clamp_log_mean_mul_neg_pow_rsub_sum_1[grid(1)](
            arg0_1, arg1_1, buf9, 1, 64, XBLOCK=1, num_warps=2, num_stages=1)
        del arg0_1
        del arg1_1
        buf10 = buf2
        del buf2
        buf11 = buf10
        del buf10
        triton_per_fused_add_div_dot_lift_fresh_mean_mul_relu_rsub_sub_2[grid
            (1)](buf11, buf0, buf3, buf4, buf5, buf6, buf7, buf9, 1, 256,
            num_warps=2, num_stages=1)
        del buf0
        del buf3
        del buf4
        del buf5
        del buf6
        del buf7
        del buf9
    return buf11,


class LovaszHingeLoss(nn.Module):
    """
    This class implements the lovasz hinge loss which is the continuous of the IoU for binary segmentation.
    Source: https://github.com/bermanmaxim/LovaszSoftmax
    """

    def __init__(self) ->None:
        """
        Constructor method
        """
        super(LovaszHingeLoss, self).__init__()

    def _calc_grad(self, label_sorted: 'torch.Tensor') ->torch.Tensor:
        """
        Method computes the gradients of the sorted and flattened label
        :param label_sorted: (torch.Tensor) Sorted and flattened label of shape [n]
        :return: (torch.Tensor) Gradient tensor
        """
        label_sum = label_sorted.sum()
        intersection = label_sum - label_sorted.cumsum(dim=0)
        union = label_sum + (1 - label_sorted).cumsum(dim=0)
        iou = 1.0 - intersection / union
        iou[1:] = iou[1:] - iou[0:-1]
        return iou

    def forward(self, prediction: 'torch.Tensor', label: 'torch.Tensor'
        ) ->torch.Tensor:
        """
        Forward pass computes the dice loss
        :param prediction: (torch.Tensor) Prediction
        :param label: (torch.Tensor) Label
        :return: (torch.Tensor) Dice loss value
        """
        prediction = prediction.flatten(start_dim=0)
        label = label.flatten(start_dim=0)
        signs = 2.0 * label - 1.0
        error = 1.0 - prediction * signs
        errors_sorted, permutation = torch.sort(error, dim=0, descending=True)
        label_sorted = label[permutation]
        grad = self._calc_grad(label_sorted)
        loss = torch.dot(F.relu(errors_sorted), grad)
        return loss


class DiceLoss(nn.Module):
    """
    This class implements the dice loss for multiple instances
    """

    def __init__(self, smooth_factor: 'float'=1.0) ->None:
        super(DiceLoss, self).__init__()
        self.smooth_factor = smooth_factor

    def __repr__(self):
        """
        Get representation of the loss module
        :return: (str) String including information
        """
        return '{}, smooth factor={}'.format(self.__class__.__name__, self.
            smooth_factor)

    def forward(self, prediction: 'torch.Tensor', label: 'torch.Tensor'
        ) ->torch.Tensor:
        """
        Forward pass computes the dice loss
        :param prediction: (torch.Tensor) Prediction
        :param label: (torch.Tensor) Label
        :return: (torch.Tensor) Dice loss value
        """
        prediction = prediction.flatten(start_dim=0)
        label = label.flatten(start_dim=0)
        loss = torch.tensor(1.0, dtype=torch.float32, device=prediction.device
            ) - (2.0 * torch.sum(torch.mul(prediction, label)) + self.
            smooth_factor) / (torch.sum(prediction) + torch.sum(label) +
            self.smooth_factor)
        return loss


class FocalLoss(nn.Module):
    """
    This class implements the segmentation focal loss.
    https://arxiv.org/abs/1708.02002
    """

    def __init__(self, alpha: 'float'=0.25, gamma: 'float'=2.0) ->None:
        """
        Constructor method
        :param alpha: (float) Alpha constant
        :param gamma: (float) Gamma constant (see paper)
        """
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma

    def __repr__(self):
        """
        Get representation of the loss module
        :return: (str) String including information
        """
        return '{}, alpha={}, gamma={}'.format(self.__class__.__name__,
            self.alpha, self.gamma)

    def forward(self, prediction: 'torch.Tensor', label: 'torch.Tensor'
        ) ->torch.Tensor:
        """
        Forward pass computes the binary cross entropy loss of segmentation masks
        :param prediction: (torch.Tensor) Prediction probability
        :param label: (torch.Tensor) Label one-hot encoded
        :return: (torch.Tensor) Loss value
        """
        binary_cross_entropy_loss = -(label * torch.log(prediction.clamp(
            min=1e-12)) + (1.0 - label) * torch.log((1.0 - prediction).
            clamp(min=1e-12)))
        focal_factor = prediction * label + (1.0 - prediction) * (1.0 - label)
        loss = ((1.0 - focal_factor) ** self.gamma *
            binary_cross_entropy_loss * self.alpha).sum(dim=1).mean()
        return loss


class SegmentationLossNew(nn.Module):
    """
    This class implement the segmentation loss.
    """

    def __init__(self, dice_loss: 'nn.Module'=DiceLoss(), focal_loss:
        'nn.Module'=FocalLoss(), lovasz_hinge_loss: 'nn.Module'=
        LovaszHingeLoss(), w_dice: 'float'=1.0, w_focal: 'float'=0.2,
        w_lovasz_hinge: 'float'=0.0) ->None:
        super(SegmentationLossNew, self).__init__()
        self.dice_loss = dice_loss
        self.focal_loss = focal_loss
        self.lovasz_hinge_loss = lovasz_hinge_loss
        self.w_dice = w_dice
        self.w_focal = w_focal
        self.w_lovasz_hinge = w_lovasz_hinge

    def __repr__(self):
        """
        Get representation of the loss module
        :return: (str) String including information
        """
        return ('{}, {}, w_focal={}, {}, w_dice={}, {}, w_lovasz_hinge={}'.
            format(self.__class__.__name__, self.dice_loss.__class__.
            __name__, self.w_dice, self.focal_loss.__class__.__name__, self
            .w_focal, self.lovasz_hinge_loss.__class__.__name__, self.
            w_lovasz_hinge))

    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.014882