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

# -------------------------------------------------------------
# 常量定义
# -------------------------------------------------------------
N_BATCH = 128
N_FEATURES = 512

# 损失函数参数
FULL = False
EPS = 1e-6
REDUCTION = 'mean'


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

class Model(nn.Module):
    """
    nn.GaussianNLLLoss 的纯 PyTorch 基准实现
    """

    def __init__(self, full=False, eps=1e-6, reduction='mean'):
        super().__init__()
        self.full = full
        self.eps = eps
        self.reduction = reduction

        if self.full:
            self.const_term = 0.5 * math.log(2 * math.pi)
        else:
            self.const_term = 0.0

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

        # 1. 确保 var > eps。
        #    torch.clamp(min=...) 等价于 max(var, eps)
        var_clamped = torch.clamp(var, min=self.eps)

        # 2. 计算两个主要项
        term1_log = torch.log(var_clamped)
        term2_sq_err = (input - target).pow(2) / var_clamped

        # 3. 组合
        # (N, *) 形状
        loss_unreduced = 0.5 * (term1_log + term2_sq_err) + self.const_term

        # 4. 应用 Reduciton
        if self.reduction == 'mean':
            return loss_unreduced.mean()
        elif self.reduction == 'sum':
            return loss_unreduced.sum()
        else:  # 'none'
            return loss_unreduced


def get_inputs():
    """
    生成 (N, D) 形状的输入
    """
    input = torch.randn(N_BATCH, N_FEATURES, dtype=torch.float32)
    target = torch.randn(N_BATCH, N_FEATURES, dtype=torch.float32)

    # Var 必须是正数
    var = torch.rand(N_BATCH, N_FEATURES, dtype=torch.float32) + EPS

    return [input, target, var]


def get_init_inputs():
    return [FULL, EPS, REDUCTION]

