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

# -------------------------------------------------------------
# 常量定义
# -------------------------------------------------------------
BATCH_SIZE = 32
CHANNELS = 64
WIDTH = 128  # W_in
PADDING = (3, 1)  # (padding_left, padding_right)


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

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

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

        if isinstance(padding, int):
            # F.pad 需要 (left, right) 格式
            self.padding_tuple = (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, pad_dim_1_left, ...)
        # 因为我们只 pad 最后一个维度 (dim -1)，所以元组是 (pad_L, pad_R)
        return F.pad(x, self.padding_tuple, mode='reflect')


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


def get_init_inputs():
    return [PADDING]
