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

BATCH_SIZE = 8
CHANNELS = 16
DEPTH = 16
HEIGHT = 16
WIDTH = 16

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

BLOCK_DIM_X = 8
BLOCK_DIM_Y = 8
BLOCK_DIM_Z = 8


class ModelNew(nn.Module):

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

        if isinstance(padding, int):
            self.pad_L, self.pad_R, self.pad_T, self.pad_B, self.pad_F, self.pad_K = (padding,) * 6
        else:
            self.pad_L, self.pad_R, self.pad_T, self.pad_B, self.pad_F, self.pad_K = padding

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

    def _compile_cuda_kernel(self):

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

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

        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}
        #define BLOCK_DIM_Z {self.block_dim_z}


        __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_pad3d_fused_kernel(
            const float* __restrict__ input_data,
            float* __restrict__ output_data,
            int N, int C, 
            int D_in, int H_in, int W_in, 
            int D_out, int H_out, int W_out,
            int pad_L, int pad_R,
            int pad_T, int pad_B,
            int pad_F, int pad_K
        ) {{
            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 int tid_z = threadIdx.z;


            const int64_t H_in_stride = W_in;
            const int64_t D_in_stride = H_in * W_in;

            const int64_t C_in_stride = D_in * D_in_stride; 


            const int64_t H_out_stride = W_out;
            const int64_t D_out_stride = H_out * W_out;

            const int64_t C_out_stride = D_out * D_out_stride; 


            const float* p_in_base = input_data + (n_idx * C + c_idx) * C_in_stride;
            float* p_out_base = output_data + (n_idx * C + c_idx) * C_out_stride;


            for (int k = tid_z; k < D_in; k += BLOCK_DIM_Z) {{
                for (int i = tid_y; i < H_in; i += BLOCK_DIM_Y) {{
                    for (int j = tid_x; j < W_in; j += BLOCK_DIM_X) {{
                        // (k, i, j) -> 1D index
                        int64_t in_idx = k*D_in_stride + i*H_in_stride + j;
                        s_in[in_idx] = p_in_base[in_idx];
                    }}
                }}
            }}
            __syncthreads(); // 确保 s_in 加载完成


            for (int k = tid_z; k < D_out; k += BLOCK_DIM_Z) {{
                int in_k = reflect_idx(k, pad_F, D_in);

                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);

                        // 从共享内存读取 (使用 IN 步长)
                        int64_t s_in_idx = in_k*D_in_stride + in_i*H_in_stride + in_j;

                        // 写入全局内存 (使用 OUT 步长)
                        int64_t p_out_idx = k*D_out_stride + i*H_out_stride + j;

                        p_out_base[p_out_idx] = s_in[s_in_idx];
                    }}
                }}
            }}
        }}

        // C++ 封装函数
        torch::Tensor reflection_pad3d_forward_cuda(
            torch::Tensor input,
            int pad_L, int pad_R,
            int pad_T, int pad_B,
            int pad_F, int pad_K
        ) {{
            TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
            TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
            TORCH_CHECK(input.dim() == 5, "input must be 5D (N, C, D, H, W)");

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

            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");
            TORCH_CHECK(pad_F < D_in_64, "pad_F error");
            TORCH_CHECK(pad_K < D_in_64, "pad_K error");

            const int64_t D_out_64 = D_in_64 + pad_F + pad_K;
            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, D_out_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, BLOCK_DIM_Z);


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

            reflection_pad3d_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>(D_in_64), static_cast<int>(H_in_64), static_cast<int>(W_in_64),
                static_cast<int>(D_out_64), static_cast<int>(H_out_64), static_cast<int>(W_out_64),
                pad_L, pad_R,
                pad_T, pad_B,
                pad_F, pad_K
            );

            return output;
        }}
        """

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

        self.pad_op = load_inline(
            name="reflection_pad3d_op_v2_fixed",
            cpp_sources=cpp_header,
            cuda_sources=cuda_source,
            functions=["reflection_pad3d_forward_cuda"],
            extra_cuda_cflags=nvcc_flags,
            verbose=False
        )

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

        x_cont = x.contiguous()

        return self.pad_op.reflection_pad3d_forward_cuda(
            x_cont,
            self.pad_L, self.pad_R,
            self.pad_T, self.pad_B,
            self.pad_F, self.pad_K
        )