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

# 定义用于测试的假设常量
N, C, H, W = 32, 64, 56, 56


class WelschLoss(nn.Module):
    def __init__(self, reduction='mean', beta=1.0):
        super().__init__()
        self.reduction = reduction
        self.beta = float(beta)
        self.beta_sq = self.beta * self.beta
        if reduction not in ['none', 'mean', 'sum']:
            raise ValueError("Invalid reduction mode")

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        # Calculate squared difference: (x - y)^2
        diff_sq = (input - target) ** 2

        # Calculate loss: 1 - exp(-(diff^2 / beta^2))
        loss = 1.0 - torch.exp(-diff_sq / self.beta_sq)

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:
            return loss


class Model(nn.Module):
    def __init__(self, reduction='mean', beta=1.0):
        super().__init__()
        self.op = WelschLoss(reduction, beta)

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        # 适应 benchmark 输入格式 (如果输入是列表)
        if isinstance(input, (list, tuple)): input = input[0]
        if isinstance(target, (list, tuple)): target = target[0]

        return self.op(input, target)


def get_inputs():
    input = torch.randn(N, C, H, W, dtype=torch.float32)
    target = torch.randn(N, C, H, W, dtype=torch.float32)
    return [input, target]


def get_init_inputs():
    return ['mean', 1.0]