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

BATCH_SIZE = 4096
FEATURE_DIM = 512

# --- 损失函数的参数 ---
# log_input=True: loss = exp(input) - target * input
# log_input=False: loss = input - target * log(input + eps)
LOG_INPUT = True
# full=True: 添加 Stirling's approximation
FULL = False
EPS = 1e-8


class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.criterion = nn.PoissonNLLLoss(
            log_input=LOG_INPUT,
            full=FULL,
            eps=EPS,
            reduction='mean'
        )

    def forward(self, input_tensor: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        # 在 PyTorch 中，input_tensor 是文档中的 'input'
        return self.criterion(input_tensor, target)


def get_inputs():
    # Input (log_input=True 时) 可以是任意实数
    input_tensor = torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32)

    # Target 在 Poisson 分布中代表计数，且在 'full' 模式下会计算 log(target)
    # 因此 target 必须是 >= 0 的。我们使用 rand 来确保
    target = torch.rand(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32) * 10  # 乘以 10 以便有一些 > 1

    return [input_tensor, target]


def get_init_inputs():
    return []