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 warnings
from typing import Callable
from typing import Union
from typing import Optional
from enum import Enum
import torch.nn as nn
from torch.nn.modules.loss import _Loss
import torch.multiprocessing
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__log_softmax_mul_sum_0(in_ptr0, in_ptr1, out_ptr0,
    out_ptr1, out_ptr2, out_ptr3, xnumel, rnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r1 = rindex
    x0 = xindex
    x3 = xindex // 4
    tmp0 = tl.load(in_ptr0 + (r1 + 16 * x0), xmask, other=0.0)
    tmp1 = tl.load(in_ptr1 + (r1 + 16 * x0), xmask, other=0.0)
    tmp15 = tl.load(in_ptr1 + (r1 + 64 * x3), xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp16 = tl.load(in_ptr1 + (16 + r1 + 64 * x3), xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp18 = tl.load(in_ptr1 + (32 + r1 + 64 * x3), xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp20 = tl.load(in_ptr1 + (48 + r1 + 64 * x3), xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp2 = tmp0 * tmp1
    tmp3 = tl.broadcast_to(tmp2, [XBLOCK, RBLOCK])
    tmp5 = tl.where(xmask, tmp3, 0)
    tmp6 = tl.sum(tmp5, 1)[:, None]
    tmp7 = tl.broadcast_to(tmp0, [XBLOCK, RBLOCK])
    tmp9 = tl.where(xmask, tmp7, 0)
    tmp10 = tl.sum(tmp9, 1)[:, None]
    tmp11 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp13 = tl.where(xmask, tmp11, 0)
    tmp14 = tl.sum(tmp13, 1)[:, None]
    tmp17 = triton_helpers.maximum(tmp15, tmp16)
    tmp19 = triton_helpers.maximum(tmp17, tmp18)
    tmp21 = triton_helpers.maximum(tmp19, tmp20)
    tmp22 = tmp1 - tmp21
    tl.store(out_ptr3 + (r1 + 16 * x0), tmp22, xmask)
    tl.store(out_ptr0 + x0, tmp6, xmask)
    tl.store(out_ptr1 + x0, tmp10, xmask)
    tl.store(out_ptr2 + x0, tmp14, xmask)


@triton.jit
def triton_per_fused_add_div_mean_mul_neg_1(in_ptr0, in_ptr1, in_ptr2,
    out_ptr0, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 16
    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
    tmp0 = tl.load(in_ptr0 + r0, None)
    tmp6 = tl.load(in_ptr1 + r0, None)
    tmp7 = tl.load(in_ptr2 + r0, None)
    tmp1 = 2.0
    tmp2 = tmp0 * tmp1
    tmp3 = 1e-05
    tmp4 = tmp2 + tmp3
    tmp5 = -tmp4
    tmp8 = tmp6 + tmp7
    tmp9 = tmp8 + tmp3
    tmp10 = tmp5 / tmp9
    tmp11 = tl.broadcast_to(tmp10, [XBLOCK, RBLOCK])
    tmp13 = tl.sum(tmp11, 1)[:, None]
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp13, None)


@triton.jit
def triton_per_fused_add_argmax_div_mean_mul_neg_nll_loss2d_forward_2(
    in_out_ptr0, 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 + (r0 + 64 * r1), None)
    tmp1 = tl.load(in_ptr0 + (16 + r0 + 64 * r1), None)
    tmp17 = tl.load(in_ptr0 + (32 + r0 + 64 * r1), None)
    tmp32 = tl.load(in_ptr0 + (48 + r0 + 64 * r1), None)
    tmp56 = tl.load(in_ptr1 + (r0 + 64 * r1), None)
    tmp58 = tl.load(in_ptr1 + (16 + r0 + 64 * r1), None)
    tmp61 = tl.load(in_ptr1 + (32 + r0 + 64 * r1), None)
    tmp64 = tl.load(in_ptr1 + (48 + r0 + 64 * r1), None)
    tmp79 = tl.load(in_out_ptr0 + 0)
    tmp80 = tl.broadcast_to(tmp79, [XBLOCK, 1])
    tmp2 = tmp0 > tmp1
    tmp3 = tmp0 == tmp1
    tmp4 = tmp0 != tmp0
    tmp5 = tmp1 != tmp1
    tmp6 = tmp4 > tmp5
    tmp7 = tmp2 | tmp6
    tmp8 = tmp4 & tmp5
    tmp9 = tmp3 | tmp8
    tmp10 = tl.full([1, 1], 0, tl.int64)
    tmp11 = tl.full([1, 1], 1, tl.int64)
    tmp12 = tmp10 < tmp11
    tmp13 = tmp9 & tmp12
    tmp14 = tmp7 | tmp13
    tmp15 = tl.where(tmp14, tmp0, tmp1)
    tmp16 = tl.where(tmp14, tmp10, tmp11)
    tmp18 = tmp15 > tmp17
    tmp19 = tmp15 == tmp17
    tmp20 = tmp15 != tmp15
    tmp21 = tmp17 != tmp17
    tmp22 = tmp20 > tmp21
    tmp23 = tmp18 | tmp22
    tmp24 = tmp20 & tmp21
    tmp25 = tmp19 | tmp24
    tmp26 = tl.full([1, 1], 2, tl.int64)
    tmp27 = tmp16 < tmp26
    tmp28 = tmp25 & tmp27
    tmp29 = tmp23 | tmp28
    tmp30 = tl.where(tmp29, tmp15, tmp17)
    tmp31 = tl.where(tmp29, tmp16, tmp26)
    tmp33 = tmp30 > tmp32
    tmp34 = tmp30 == tmp32
    tmp35 = tmp30 != tmp30
    tmp36 = tmp32 != tmp32
    tmp37 = tmp35 > tmp36
    tmp38 = tmp33 | tmp37
    tmp39 = tmp35 & tmp36
    tmp40 = tmp34 | tmp39
    tmp41 = tl.full([1, 1], 3, tl.int64)
    tmp42 = tmp31 < tmp41
    tmp43 = tmp40 & tmp42
    tmp44 = tmp38 | tmp43
    tl.where(tmp44, tmp30, tmp32)
    tmp46 = tl.where(tmp44, tmp31, tmp41)
    tmp47 = tl.full([1, 1], -100, tl.int64)
    tmp48 = tmp46 != tmp47
    tmp49 = tl.where(tmp48, tmp46, tmp10)
    tmp50 = tl.full([XBLOCK, RBLOCK], 4, tl.int32)
    tmp51 = tmp49 + tmp50
    tmp52 = tmp49 < 0
    tmp53 = tl.where(tmp52, tmp51, tmp49)
    tl.device_assert((0 <= tmp53) & (tmp53 < 4),
        'index out of bounds: 0 <= tmp53 < 4')
    tmp55 = tl.load(in_ptr1 + (r0 + 16 * tmp53 + 64 * r1), None)
    tmp57 = tl_math.exp(tmp56)
    tmp59 = tl_math.exp(tmp58)
    tmp60 = tmp57 + tmp59
    tmp62 = tl_math.exp(tmp61)
    tmp63 = tmp60 + tmp62
    tmp65 = tl_math.exp(tmp64)
    tmp66 = tmp63 + tmp65
    tmp67 = tl_math.log(tmp66)
    tmp68 = tmp55 - tmp67
    tmp69 = -tmp68
    tmp70 = 0.0
    tmp71 = tl.where(tmp48, tmp69, tmp70)
    tmp72 = tl.broadcast_to(tmp71, [XBLOCK, RBLOCK])
    tmp74 = tl.sum(tmp72, 1)[:, None]
    tmp75 = tmp48.to(tl.int64)
    tmp76 = tl.broadcast_to(tmp75, [XBLOCK, RBLOCK])
    tmp78 = tl.sum(tmp76, 1)[:, None]
    tmp81 = 16.0
    tmp82 = tmp80 / tmp81
    tmp83 = tmp78.to(tl.float32)
    tmp84 = tmp74 / tmp83
    tmp85 = tmp82 + tmp84
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp85, 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((4, 4), (4, 1), torch.float32)
        buf1 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf2 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf5 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_per_fused__log_softmax_mul_sum_0[grid(16)](arg1_1, arg0_1,
            buf0, buf1, buf2, buf5, 16, 16, XBLOCK=1, num_warps=2, num_stages=1
            )
        del arg0_1
        buf3 = empty_strided_cuda((), (), torch.float32)
        triton_per_fused_add_div_mean_mul_neg_1[grid(1)](buf0, buf1, buf2,
            buf3, 1, 16, XBLOCK=1, num_warps=2, num_stages=1)
        del buf0
        del buf1
        del buf2
        buf8 = buf3
        del buf3
        triton_per_fused_add_argmax_div_mean_mul_neg_nll_loss2d_forward_2[grid
            (1)](buf8, arg1_1, buf5, 1, 64, XBLOCK=1, num_warps=2, num_stages=1
            )
        del arg1_1
        del buf5
    return buf8,


class LossReduction(Enum):
    """
    See also:
        - :py:class:`monai.losses.dice.DiceLoss`
        - :py:class:`monai.losses.dice.GeneralizedDiceLoss`
        - :py:class:`monai.losses.focal_loss.FocalLoss`
        - :py:class:`monai.losses.tversky.TverskyLoss`
    """
    NONE = 'none'
    MEAN = 'mean'
    SUM = 'sum'
    SUM_MEAN = 'sum_mean'


class DiceLoss(_Loss):
    """
    Compute average Dice loss between two tensors. It can support both multi-classes and multi-labels tasks.
    Input logits `input` (BNHW[D] where N is number of classes) is compared with ground truth `target` (BNHW[D]).
    Axis N of `input` is expected to have logit predictions for each class rather than being image channels,
    while the same axis of `target` can be 1 or N (one-hot format). The `smooth_nr` and `smooth_dr` parameters are
    values added to the intersection and union components of the inter-over-union calculation to smooth results
    respectively, these values should be small. The `include_background` class attribute can be set to False for
    an instance of DiceLoss to exclude the first category (channel index 0) which is by convention assumed to be
    background. If the non-background segmentations are small compared to the total image size they can get
    overwhelmed by the signal from the background so excluding it in such cases helps convergence.

    Milletari, F. et. al. (2016) V-Net: Fully Convolutional Neural Networks forVolumetric Medical Image Segmentation, 3DV, 2016.

    """

    def __init__(self, include_background: 'bool'=True, to_onehot_y: 'bool'
        =False, sigmoid: 'bool'=False, softmax: 'bool'=False, other_act:
        'Optional[Callable]'=None, squared_pred: 'bool'=False, jaccard:
        'bool'=False, reduction: 'Union[LossReduction, str]'=LossReduction.
        MEAN, smooth_nr: 'float'=1e-05, smooth_dr: 'float'=1e-05, batch:
        'bool'=False, my_dice: 'bool'=False) ->None:
        """
        Args:
            include_background: if False channel index 0 (background category) is excluded from the calculation.
            to_onehot_y: whether to convert `y` into the one-hot format. Defaults to False.
            sigmoid: if True, apply a sigmoid function to the prediction.
            softmax: if True, apply a softmax function to the prediction.
            other_act: if don't want to use `sigmoid` or `softmax`, use other callable function to execute
                other activation layers, Defaults to ``None``. for example:
                `other_act = torch.tanh`.
            squared_pred: use squared versions of targets and predictions in the denominator or not.
            jaccard: compute Jaccard Index (soft IoU) instead of dice or not.
            reduction: {``"none"``, ``"mean"``, ``"sum"``}
                Specifies the reduction to apply to the output. Defaults to ``"mean"``.

                - ``"none"``: no reduction will be applied.
                - ``"mean"``: the sum of the output will be divided by the number of elements in the output.
                - ``"sum"``: the output will be summed.

            smooth_nr: a small constant added to the numerator to avoid zero.
            smooth_dr: a small constant added to the denominator to avoid nan.
            batch: whether to sum the intersection and union areas over the batch dimension before the dividing.
                Defaults to False, a Dice loss value is computed independently from each item in the batch
                before any `reduction`.

        Raises:
            TypeError: When ``other_act`` is not an ``Optional[Callable]``.
            ValueError: When more than 1 of [``sigmoid=True``, ``softmax=True``, ``other_act is not None``].
                Incompatible values.

        """
        super().__init__(reduction=LossReduction(reduction).value)
        if other_act is not None and not callable(other_act):
            raise TypeError(
                f'other_act must be None or callable but is {type(other_act).__name__}.'
                )
        if int(sigmoid) + int(softmax) + int(other_act is not None) > 1:
            raise ValueError(
                'Incompatible values: more than 1 of [sigmoid=True, softmax=True, other_act is not None].'
                )
        self.include_background = include_background
        self.to_onehot_y = to_onehot_y
        self.sigmoid = sigmoid
        self.softmax = softmax
        self.other_act = other_act
        self.squared_pred = squared_pred
        self.jaccard = jaccard
        self.smooth_nr = float(smooth_nr)
        self.smooth_dr = float(smooth_dr)
        self.batch = batch
        self.my_dice = my_dice

    def forward(self, input: 'torch.Tensor', target: 'torch.Tensor'
        ) ->torch.Tensor:
        """
        Args:
            input: the shape should be BNH[WD].
            target: the shape should be BNH[WD].

        Raises:
            ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"].

        """
        if self.sigmoid:
            input = torch.sigmoid(input)
        n_pred_ch = input.shape[1]
        if self.softmax:
            if n_pred_ch == 1:
                warnings.warn(
                    'single channel prediction, `softmax=True` ignored.')
            else:
                input = torch.softmax(input, 1)
        if self.other_act is not None:
            input = self.other_act(input)
        if self.to_onehot_y:
            if n_pred_ch == 1:
                warnings.warn(
                    'single channel prediction, `to_onehot_y=True` ignored.')
            else:
                target = one_hot(target, num_classes=n_pred_ch)
        assert target.shape == input.shape, f'ground truth has differing shape ({target.shape}) from input ({input.shape})'
        reduce_axis = list(range(2, len(input.shape)))
        if self.batch:
            reduce_axis = [0] + reduce_axis
        intersection = torch.sum(target * input, dim=reduce_axis)
        if self.squared_pred:
            target = torch.pow(target, 2)
            input = torch.pow(input, 2)
        ground_o = torch.sum(target, dim=reduce_axis)
        pred_o = torch.sum(input, dim=reduce_axis)
        denominator = ground_o + pred_o
        if self.jaccard:
            denominator = 2.0 * (denominator - intersection)
        if self.my_dice:

            def count_nonzeros(input, dim):
                mask = input != 0.0
                return torch.sum(mask, dim)
            f: 'torch.Tensor' = 2.0 * intersection / (denominator + self.
                smooth_nr)
            if not self.include_background:
                if self.batch:
                    f = f[1:]
                    ground_o = ground_o[1:]
                else:
                    f = f[:, 1:]
                    ground_o = ground_o[:, 1:]
            f = -(torch.sum(f, dim=-1) / count_nonzeros(ground_o, dim=-1))
        else:
            f: 'torch.Tensor' = -(2.0 * intersection + self.smooth_nr) / (
                denominator + self.smooth_dr)
            if not self.include_background:
                if self.batch:
                    f = f[1:]
                else:
                    f = f[:, 1:]
        if self.reduction == LossReduction.MEAN.value:
            f = torch.mean(f)
        elif self.reduction == LossReduction.SUM.value:
            f = torch.sum(f)
        elif self.reduction == LossReduction.NONE.value:
            pass
        else:
            raise ValueError(
                f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].'
                )
        return f


class DiceCELossNew(_Loss):
    """
    Compute both Dice loss and Cross Entropy Loss, and return the sum of these two losses.
    Input logits `input` (BNHW[D] where N is number of classes) is compared with ground truth `target` (BNHW[D]).
    Axis N of `input` is expected to have logit predictions for each class rather than being image channels,
    while the same axis of `target` can be 1 or N (one-hot format). The `smooth_nr` and `smooth_dr` parameters are
    values added for dice loss part to the intersection and union components of the inter-over-union calculation
    to smooth results respectively, these values should be small. The `include_background` class attribute can be
    set to False for an instance of the loss to exclude the first category (channel index 0) which is by convention
    assumed to be background. If the non-background segmentations are small compared to the total image size they can get
    overwhelmed by the signal from the background so excluding it in such cases helps convergence.
    """

    def __init__(self, include_background: 'bool'=True, to_onehot_y: 'bool'
        =False, sigmoid: 'bool'=False, softmax: 'bool'=False, other_act:
        'Optional[Callable]'=None, squared_pred: 'bool'=False, jaccard:
        'bool'=False, reduction: 'str'='mean', smooth_nr: 'float'=1e-05,
        smooth_dr: 'float'=1e-05, batch: 'bool'=False, ce_weight:
        'Optional[torch.Tensor]'=None, negate_dice: 'bool'=False, my_dice:
        'bool'=False) ->None:
        """
        Args:
            ``ce_weight`` is only used for cross entropy loss, ``reduction`` is used for both losses and other
            parameters are only used for dice loss.
            include_background: if False channel index 0 (background category) is excluded from the calculation.
            to_onehot_y: whether to convert `y` into the one-hot format. Defaults to False.
            sigmoid: if True, apply a sigmoid function to the prediction.
            softmax: if True, apply a softmax function to the prediction.
            other_act: if don't want to use `sigmoid` or `softmax`, use other callable function to execute
                other activation layers, Defaults to ``None``. for example:
                `other_act = torch.tanh`.
            squared_pred: use squared versions of targets and predictions in the denominator or not.
            jaccard: compute Jaccard Index (soft IoU) instead of dice or not.
            reduction: {``"mean"``, ``"sum"``}
                Specifies the reduction to apply to the output. Defaults to ``"mean"``. The dice loss should
                as least reduce the spatial dimensions, which is different from cross entropy loss, thus here
                the ``none`` option cannot be used.
                - ``"mean"``: the sum of the output will be divided by the number of elements in the output.
                - ``"sum"``: the output will be summed.
            smooth_nr: a small constant added to the numerator to avoid zero.
            smooth_dr: a small constant added to the denominator to avoid nan.
            batch: whether to sum the intersection and union areas over the batch dimension before the dividing.
                Defaults to False, a Dice loss value is computed independently from each item in the batch
                before any `reduction`.
            ce_weight: a rescaling weight given to each class for cross entropy loss.
                See ``torch.nn.CrossEntropyLoss()`` for more information.
        """
        super().__init__()
        self.dice = DiceLoss(include_background=include_background,
            to_onehot_y=to_onehot_y, sigmoid=sigmoid, softmax=softmax,
            other_act=other_act, squared_pred=squared_pred, jaccard=jaccard,
            reduction=reduction, smooth_nr=smooth_nr, smooth_dr=smooth_dr,
            batch=batch, my_dice=my_dice)
        self.cross_entropy = nn.CrossEntropyLoss(weight=ce_weight,
            reduction=reduction)
        self.negate_dice = negate_dice

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