import torch
import torch.nn as nn
import torch.nn.functional as F

# -------------------------------------------------------------
# 常量定义
# -------------------------------------------------------------
N, C, H, W = 8, 10, 16, 16  # (N, C, H, W)

# 损失函数参数
WEIGHT = torch.rand(C, dtype=torch.float32)  # (C,)
IGNORE_INDEX = -100
REDUCTION = 'mean'


# -------------------------------------------------------------

class Model(nn.Module):
    """
    nn.NLLLoss 的纯 PyTorch 基准实现
    (K-dim, 2D-example)
    """

    def __init__(self, weight=None, size_average=None, ignore_index=-100,
                 reduce=None, reduction='mean'):
        super().__init__()

        # 处理已弃用的 size_average 和 reduce
        if size_average is not None or reduce is not None:
            # (省略... 遵循 torch.nn.modules.loss)
            pass

        self.reduction = reduction
        self.ignore_index = ignore_index

        # 确保 weight 在正确的设备上
        if weight is not None:
            self.register_buffer('weight', weight)
        else:
            self.weight = None

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:

        input_flat = input.view(N, C, -1)

        target_flat = target.view(N, -1)


        loss_unreduced = input_flat.gather(dim=1, index=target_flat.unsqueeze(1))
        loss_unreduced = -loss_unreduced.squeeze(1)  # (N, H*W)


        if self.weight is not None:

            weights_applied = self.weight[target_flat]
            loss_unreduced = loss_unreduced * weights_applied
        else:

            weights_applied = torch.ones_like(target_flat, dtype=input.dtype)


        mask = (target_flat != self.ignore_index)
        loss_unreduced = loss_unreduced * mask
        weights_applied = weights_applied * mask


        if self.reduction == 'mean':

            total_weight = weights_applied.sum()
            if total_weight == 0:
                return torch.tensor(0.0, device=input.device, dtype=input.dtype)
            return loss_unreduced.sum() / total_weight

        elif self.reduction == 'sum':
            return loss_unreduced.sum()

        else:
            return loss_unreduced.view_as(target)


def get_inputs():

    input_log_probs = F.log_softmax(torch.randn(N, C, H, W, dtype=torch.float32), dim=1)
    target = torch.empty(N, H, W, dtype=torch.long).random_(0, C)


    target.view(-1)[::10] = IGNORE_INDEX

    return [input_log_probs, target]


def get_init_inputs():

    return [WEIGHT, None, IGNORE_INDEX, None, REDUCTION]
