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

# -------------------------------------------------------------
# 常量定义
# -------------------------------------------------------------
BATCH_SIZE = 32
CHANNELS = 64
HEIGHT = 32  # H_in
WIDTH = 32  # W_in

# (pad_L, pad_R, pad_T, pad_B)
PADDING = (1, 1, 2, 0)


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

class Model(nn.Module):
    """
    nn.ReflectionPad2d 的纯 PyTorch 基准实现
    (使用 F.pad)
    """

    def __init__(self, padding):
        super().__init__()

        if isinstance(padding, int):
            # F.pad 需要 (left, right, top, bottom) 格式
            self.padding_tuple = (padding, padding, padding, padding)
        else:
            self.padding_tuple = padding

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # F.pad 的 padding 格式是 (pad_dim_0_left, pad_dim_0_right, ...)
        # 对应 (N, C, H, W)，我们需要 pad 最后两个维度
        # F.pad 接受的顺序是 (pad_W_left, pad_W_right, pad_H_top, pad_H_bottom)
        return F.pad(x, self.padding_tuple, mode='reflect')


def get_inputs():
    """
    生成一个 (N, C, H, W) 形状的输入
    """
    x = torch.randn(BATCH_SIZE, CHANNELS, HEIGHT, WIDTH, dtype=torch.float32)
    return [x]


def get_init_inputs():
    return [PADDING]
