import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# -------------------------------------------------------------
# 常量定义
# -------------------------------------------------------------
BATCH_SIZE = 32
CHANNELS = 64
WIDTH = 128  # W_in
PADDING = (3, 1)  # (padding_left, padding_right)
BLOCK_SIZE = 256  # CUDA Block 维度


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

class ModelNew(nn.Module):
    """
    ReflectionPad1d 的高性能 CUDA 融合核函数实现
    (修复了编译错误)
    """

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

        if isinstance(padding, int):
            self.pad_L = padding
            self.pad_R = padding
        else:
            self.pad_L = padding[0]
            self.pad_R = padding[1]

        self.block_size = BLOCK_SIZE
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):

        cpp_header = f"""
        #include <torch/extension.h>

        // C++ 接口
        torch::Tensor reflection_pad1d_forward_cuda(
            torch::Tensor input,
            int pad_L,
            int pad_R
        );
        """

        cuda_source = f"""
        #include <torch/extension.h> 
        #include <cuda_runtime.h>

        // [修复] 将 #define 移至此处
        #define BLOCK_SIZE {self.block_size}

        /*
         * ReflectionPad1d 融合核函数
         */
        __global__ void reflection_pad1d_fused_kernel(
            const float* __restrict__ input_data,
            float* __restrict__ output_data,
            int N, int C, int W_in, int W_out,
            int pad_L, int pad_R
        ) {{ // <-- f-string 转义
            const int n_idx = blockIdx.x; 
            const int c_idx = blockIdx.y; 
            const int tid = threadIdx.x;

            const float* p_in = input_data + (n_idx * C + c_idx) * W_in;
            float* p_out = output_data + (n_idx * C + c_idx) * W_out;

            // [修复] BLOCK_SIZE 现在可见
            for (int j = tid; j < W_out; j += BLOCK_SIZE) {{ // <-- f-string 转义
                int in_idx = 0;

                if (j < pad_L) {{
                    in_idx = pad_L - j;
                }} else if (j < (pad_L + W_in)) {{
                    in_idx = j - pad_L;
                }} else {{
                    int j_rel = j - (pad_L + W_in);
                    in_idx = W_in - 2 - j_rel;
                }}

                p_out[j] = p_in[in_idx];
            }}
        }} 

        // C++ 封装函数
        // [修复] torch.Tensor -> torch::Tensor
        torch::Tensor reflection_pad1d_forward_cuda(
            torch::Tensor input,
            int pad_L,
            int pad_R
        ) {{ 
            TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
            TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
            TORCH_CHECK(input.dim() == 3, "input must be 3D (N, C, W)");

            const int64_t N_64 = input.size(0);
            const int64_t C_64 = input.size(1);
            const int64_t W_in_64 = input.size(2);

            TORCH_CHECK(pad_L < W_in_64, "padding_left should be less than input width");
            TORCH_CHECK(pad_R < W_in_64, "padding_right should be less than input width");

            const int64_t W_out_64 = W_in_64 + pad_L + pad_R;

            auto output = torch::empty({{N_64, C_64, W_out_64}}, input.options());

            dim3 grid_dim(N_64, C_64); 
            dim3 block_dim(BLOCK_SIZE);

            reflection_pad1d_fused_kernel<<<grid_dim, block_dim>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                static_cast<int>(N_64),
                static_cast<int>(C_64),
                static_cast<int>(W_in_64),
                static_cast<int>(W_out_64),
                pad_L,
                pad_R
            );

            return output;
        }} 
        """

        # JIT (Just-In-Time) 编译
        self.pad_op = load_inline(
            name="reflection_pad1d_op_v3_fixed",  # 更改名称以避免缓存
            cpp_sources=cpp_header,
            cuda_sources=cuda_source,
            functions=["reflection_pad1d_forward_cuda"],
            verbose=False  # 如果还报错，请设为 True
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:

        # 调用我们编译好的 CUDA C++ 函数
        return self.pad_op.reflection_pad1d_forward_cuda(
            x,
            self.pad_L,
            self.pad_R
        )