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

# -------------------------------------------------------------
# 常量定义 (与你之前的代码一致)
# -------------------------------------------------------------
BATCH_SIZE = 32
CHANNELS = 64
HEIGHT = 32  # H_in
WIDTH = 32  # W_in
PADDING = (1, 1, 2, 0)
BLOCK_DIM_X = 16
BLOCK_DIM_Y = 16


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

class ModelNew(nn.Module):
    """
    ReflectionPad2d 的高性能 CUDA 融合核函数实现
    (V3: 修复了 'contiguous' 运行时错误)
    """

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

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

        self.block_dim_x = BLOCK_DIM_X
        self.block_dim_y = BLOCK_DIM_Y
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):

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

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

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

        #define BLOCK_DIM_X {self.block_dim_x}
        #define BLOCK_DIM_Y {self.block_dim_y}

        __device__ inline int reflect_idx(
            int j, int pad_before, int W_in
        ) {{
            if (j < pad_before) {{
                return pad_before - j;
            }} else if (j < (pad_before + W_in)) {{
                return j - pad_before;
            }} else {{
                int j_rel = j - (pad_before + W_in);
                return W_in - 2 - j_rel;
            }}
        }}

        __global__ void reflection_pad2d_fused_kernel(
            const float* __restrict__ input_data,
            float* __restrict__ output_data,
            int N, int C, 
            int H_in, int W_in, 
            int H_out, int W_out,
            int pad_L, int pad_R,
            int pad_T, int pad_B
        ) {{
            extern __shared__ float s_in[];

            const int n_idx = blockIdx.x; 
            const int c_idx = blockIdx.y; 
            const int tid_x = threadIdx.x;
            const int tid_y = threadIdx.y;

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

            // Pass 1: Load to shared memory
            for (int i = tid_y; i < H_in; i += BLOCK_DIM_Y) {{
                for (int j = tid_x; j < W_in; j += BLOCK_DIM_X) {{
                    s_in[i * W_in + j] = p_in[i * W_in + j];
                }}
            }}
            __syncthreads();

            // Pass 2: Compute and store from shared memory
            for (int i = tid_y; i < H_out; i += BLOCK_DIM_Y) {{
                int in_i = reflect_idx(i, pad_T, H_in);

                for (int j = tid_x; j < W_out; j += BLOCK_DIM_X) {{
                    int in_j = reflect_idx(j, pad_L, W_in);
                    p_out[i * W_out + j] = s_in[in_i * W_in + in_j];
                }}
            }}
        }}

        // C++ 封装函数
        torch::Tensor reflection_pad2d_forward_cuda(
            torch::Tensor input,
            int pad_L, int pad_R,
            int pad_T, int pad_B
        ) {{
            // 这个检查现在是安全的，因为我们在 Python 中确保了连续性
            TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
            TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
            TORCH_CHECK(input.dim() == 4, "input must be 4D (N, C, H, W)");

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

            TORCH_CHECK(pad_L < W_in_64, "pad_L error");
            TORCH_CHECK(pad_R < W_in_64, "pad_R error");
            TORCH_CHECK(pad_T < H_in_64, "pad_T error");
            TORCH_CHECK(pad_B < H_in_64, "pad_B error");

            const int64_t H_out_64 = H_in_64 + pad_T + pad_B;
            const int64_t W_out_64 = W_in_64 + pad_L + pad_R;

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

            dim3 grid_dim(N_64, C_64); 
            dim3 block_dim(BLOCK_DIM_X, BLOCK_DIM_Y);

            const int shared_mem_size = H_in_64 * W_in_64 * sizeof(float);

            reflection_pad2d_fused_kernel<<<grid_dim, block_dim, shared_mem_size>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                static_cast<int>(N_64), static_cast<int>(C_64),
                static_cast<int>(H_in_64), static_cast<int>(W_in_64),
                static_cast<int>(H_out_64), static_cast<int>(W_out_64),
                pad_L, pad_R,
                pad_T, pad_B
            );

            return output;
        }}
        """

        nvcc_flags = [
            '-O3',
            '--use_fast_math',
            '--expt-relaxed-constexpr'
        ]

        self.pad_op = load_inline(
            name="reflection_pad2d_op_v2_fixed",  # (与 V2 编译的二进制文件相同)
            cpp_sources=cpp_header,
            cuda_sources=cuda_source,
            functions=["reflection_pad2d_forward_cuda"],
            extra_cuda_cflags=nvcc_flags,
            verbose=False
        )

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

        # [修复]
        # 必须确保张量是连续的，才能传递给 C++/CUDA
        x_cont = x.contiguous()

        # 调用我们编译好的 CUDA C++ 函数
        return self.pad_op.reflection_pad2d_forward_cuda(
            x_cont,  # 传递连续的张量
            self.pad_L, self.pad_R,
            self.pad_T, self.pad_B
        )
        import torch


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

# -------------------------------------------------------------
# 常量定义 (与你之前的代码一致)
# -------------------------------------------------------------
BATCH_SIZE = 32
CHANNELS = 64
HEIGHT = 32  # H_in
WIDTH = 32  # W_in
PADDING = (1, 1, 2, 0)
BLOCK_DIM_X = 16
BLOCK_DIM_Y = 16


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

class ModelNew(nn.Module):
    """
    ReflectionPad2d 的高性能 CUDA 融合核函数实现
    (V3: 修复了 'contiguous' 运行时错误)
    """

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

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

        self.block_dim_x = BLOCK_DIM_X
        self.block_dim_y = BLOCK_DIM_Y
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):

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

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

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

        #define BLOCK_DIM_X {self.block_dim_x}
        #define BLOCK_DIM_Y {self.block_dim_y}

        __device__ inline int reflect_idx(
            int j, int pad_before, int W_in
        ) {{
            if (j < pad_before) {{
                return pad_before - j;
            }} else if (j < (pad_before + W_in)) {{
                return j - pad_before;
            }} else {{
                int j_rel = j - (pad_before + W_in);
                return W_in - 2 - j_rel;
            }}
        }}

        __global__ void reflection_pad2d_fused_kernel(
            const float* __restrict__ input_data,
            float* __restrict__ output_data,
            int N, int C, 
            int H_in, int W_in, 
            int H_out, int W_out,
            int pad_L, int pad_R,
            int pad_T, int pad_B
        ) {{
            extern __shared__ float s_in[];

            const int n_idx = blockIdx.x; 
            const int c_idx = blockIdx.y; 
            const int tid_x = threadIdx.x;
            const int tid_y = threadIdx.y;

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

            // Pass 1: Load to shared memory
            for (int i = tid_y; i < H_in; i += BLOCK_DIM_Y) {{
                for (int j = tid_x; j < W_in; j += BLOCK_DIM_X) {{
                    s_in[i * W_in + j] = p_in[i * W_in + j];
                }}
            }}
            __syncthreads();

            // Pass 2: Compute and store from shared memory
            for (int i = tid_y; i < H_out; i += BLOCK_DIM_Y) {{
                int in_i = reflect_idx(i, pad_T, H_in);

                for (int j = tid_x; j < W_out; j += BLOCK_DIM_X) {{
                    int in_j = reflect_idx(j, pad_L, W_in);
                    p_out[i * W_out + j] = s_in[in_i * W_in + in_j];
                }}
            }}
        }}

        // C++ 封装函数
        torch::Tensor reflection_pad2d_forward_cuda(
            torch::Tensor input,
            int pad_L, int pad_R,
            int pad_T, int pad_B
        ) {{
            // 这个检查现在是安全的，因为我们在 Python 中确保了连续性
            TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
            TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
            TORCH_CHECK(input.dim() == 4, "input must be 4D (N, C, H, W)");

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

            TORCH_CHECK(pad_L < W_in_64, "pad_L error");
            TORCH_CHECK(pad_R < W_in_64, "pad_R error");
            TORCH_CHECK(pad_T < H_in_64, "pad_T error");
            TORCH_CHECK(pad_B < H_in_64, "pad_B error");

            const int64_t H_out_64 = H_in_64 + pad_T + pad_B;
            const int64_t W_out_64 = W_in_64 + pad_L + pad_R;

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

            dim3 grid_dim(N_64, C_64); 
            dim3 block_dim(BLOCK_DIM_X, BLOCK_DIM_Y);

            const int shared_mem_size = H_in_64 * W_in_64 * sizeof(float);

            reflection_pad2d_fused_kernel<<<grid_dim, block_dim, shared_mem_size>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                static_cast<int>(N_64), static_cast<int>(C_64),
                static_cast<int>(H_in_64), static_cast<int>(W_in_64),
                static_cast<int>(H_out_64), static_cast<int>(W_out_64),
                pad_L, pad_R,
                pad_T, pad_B
            );

            return output;
        }}
        """

        nvcc_flags = [
            '-O3',
            '--use_fast_math',
            '--expt-relaxed-constexpr'
        ]

        self.pad_op = load_inline(
            name="reflection_pad2d_op_v2_fixed",  # (与 V2 编译的二进制文件相同)
            cpp_sources=cpp_header,
            cuda_sources=cuda_source,
            functions=["reflection_pad2d_forward_cuda"],
            extra_cuda_cflags=nvcc_flags,
            verbose=False
        )

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

        # [修复]
        # 必须确保张量是连续的，才能传递给 C++/CUDA
        x_cont = x.contiguous()

        # 调用我们编译好的 CUDA C++ 函数
        return self.pad_op.reflection_pad2d_forward_cuda(
            x_cont,  # 传递连续的张量
            self.pad_L, self.pad_R,
            self.pad_T, self.pad_B
        )
