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

BATCH_SIZE = 8
CHANNELS = 16
DEPTH = 16  # D_in
HEIGHT = 16  # H_in
WIDTH = 16  # W_in

PADDING = (1, 1, 2, 2, 1, 0)


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

class Model(nn.Module):

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

        if isinstance(padding, int):
            # F.pad 需要 6-tuple
            self.padding_tuple = (padding,) * 6
        else:
            self.padding_tuple = padding

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # F.pad 5D 张量 (N, C, D, H, W)
        # 填充顺序: (pad_W_L, pad_W_R, pad_H_T, pad_H_B, pad_D_F, pad_D_K)
        # 这与 nn.ReflectionPad3d 的构造函数顺序一致
        return F.pad(x, self.padding_tuple, mode='reflect')


def get_inputs():
    x = torch.randn(BATCH_SIZE, CHANNELS, DEPTH, HEIGHT, WIDTH, dtype=torch.float32)
    return [x]


def get_init_inputs():
    return [PADDING]
