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


from circularpad3d_torch import (
    BATCH_SIZE, CHANNELS, 
    D_IN, H_IN, W_IN, 
    D_OUT, H_OUT, W_OUT, 
    PADDING
)

PAD_L, PAD_R, PAD_T, PAD_B, PAD_F, PAD_K = PADDING


BLOCK_SIZE = 256
VEC_SIZE = 4

ILP_FACTOR = 4

class ModelNew(nn.Module):
    def __init__(self, padding):
        super().__init__()
        self.pad_l = padding[0]
        self.pad_r = padding[1]
        self.pad_t = padding[2]
        self.pad_b = padding[3]
        self.pad_f = padding[4]
        self.pad_k = padding[5]
        

        self._compile_cuda_kernel_static()

    def _compile_cuda_kernel_static(self):
        

        macros = f"""
        #define BLOCK_SIZE {BLOCK_SIZE}
        #define VEC_SIZE {VEC_SIZE}
        #define ILP {ILP_FACTOR}
        
        #define D_IN  {D_IN}
        #define H_IN  {H_IN}
        #define W_IN  {W_IN}
        
        #define D_OUT {D_OUT}
        #define H_OUT {H_OUT}
        #define W_OUT {W_OUT}
        
        #define PAD_L {self.pad_l}
        #define PAD_T {self.pad_t}
        #define PAD_F {self.pad_f}
        

        #define HW_IN    ({H_IN} * {W_IN})
        #define DHW_IN   ({D_IN} * {H_IN} * {W_IN})
        #define HW_OUT   ({H_OUT} * {W_OUT})
        #define DHW_OUT  ({D_OUT} * {H_OUT} * {W_OUT})
        """

        cpp_header = """
        #include <torch/extension.h>
        torch::Tensor circularpad3d_cuda(torch::Tensor input);
        """

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

        {macros}

        __device__ __forceinline__ void val_vec_set_elem(float4& v, int idx, float val) {{
            if (idx == 0) v.x = val;
            else if (idx == 1) v.y = val;
            else if (idx == 2) v.z = val;
            else v.w = val;
        }}

  
        __global__ void __launch_bounds__(BLOCK_SIZE) circularpad3d_kernel_static(
            const float* __restrict__ input_data,
            float* __restrict__ output_data,
            int total_vecs,
            int total_elements
        ) {{
            int tid = blockIdx.x * blockDim.x + threadIdx.x;
            int stride = gridDim.x * blockDim.x;

            float4* out_ptr = (float4*)output_data;

            for (int i = tid; i < total_vecs; i += stride * ILP) {{
                
                #pragma unroll
                for (int k = 0; k < ILP; ++k) {{
                    int idx = i + k * stride;
                    if (idx >= total_vecs) break;

                    int start_linear_idx = idx * VEC_SIZE;

                    int w_out_base = start_linear_idx % W_OUT;
                    int tmp = start_linear_idx / W_OUT;
     
                    int h_out_base = tmp % H_OUT;
                    tmp = tmp / H_OUT;
             
                    int d_out_base = tmp % D_OUT;
              
                    int nc_idx = tmp / D_OUT;

                    int input_base_addr = nc_idx * DHW_IN;

                    float4 val_vec;

                    #pragma unroll
                    for (int v = 0; v < VEC_SIZE; ++v) {{
                        int cur_w = w_out_base + v;
                        int cur_h = h_out_base;
                        int cur_d = d_out_base;
                        
        
                        if (cur_w >= W_OUT) {{
                            cur_w -= W_OUT;
                            cur_h++;
                            if (cur_h >= H_OUT) {{
                                cur_h = 0;
                                cur_d++;
                                if (cur_d >= D_OUT) {{
                                    cur_d = 0;
                                    input_base_addr += DHW_IN; 
                                }}
                            }}
                        }}

                        int d_offset = cur_d - PAD_F;
                        int d_in = d_offset % D_IN;
                        if (d_in < 0) d_in += D_IN;
                 
                        int h_offset = cur_h - PAD_T;
                        int h_in = h_offset % H_IN;
                        if (h_in < 0) h_in += H_IN;

                        int w_offset = cur_w - PAD_L;
                        int w_in = w_offset % W_IN;
                        if (w_in < 0) w_in += W_IN;

                 
                        int read_idx = input_base_addr 
                                      + d_in * HW_IN 
                                      + h_in * W_IN 
                                     + w_in;
                        
                        val_vec_set_elem(val_vec, v, input_data[read_idx]);
                    }}

                    out_ptr[idx] = val_vec;
                }}
            }}

            int tail_start = total_vecs * VEC_SIZE;
            for (int k = tail_start + tid; k < total_elements; k += stride) {{
                 int w_out = k % W_OUT;
                 int tmp = k / W_OUT;
                 int h_out = tmp % H_OUT;
                 tmp = tmp / H_OUT;
                 int d_out = tmp % D_OUT;
                 int nc_idx = tmp / D_OUT;
                 
                 int d_in = (d_out - PAD_F) % D_IN;
                 if (d_in < 0) d_in += D_IN;
                 
                 int h_in = (h_out - PAD_T) % H_IN;
                 if (h_in < 0) h_in += H_IN;
                 
                 int w_in = (w_out - PAD_L) % W_IN;
                 if (w_in < 0) w_in += W_IN;
                 
                 int read_idx = nc_idx * DHW_IN + d_in * HW_IN + h_in * W_IN + w_in;
                 output_data[k] = input_data[read_idx];
            }}
        }}

        torch::Tensor circularpad3d_cuda(torch::Tensor input) {{
            input = input.contiguous();
            
            auto output = torch::empty(
                {{input.size(0), input.size(1), D_OUT, H_OUT, W_OUT}}, 
                input.options()
            );
            
            int64_t total_elements = output.numel();
            int total_vecs = total_elements / VEC_SIZE;
            
            int tasks = (total_vecs + ILP - 1) / ILP;
            int grid_size = (tasks + BLOCK_SIZE - 1) / BLOCK_SIZE;
            
            if (grid_size > 65535) grid_size = 65535;
            if (grid_size == 0) grid_size = 1;

            circularpad3d_kernel_static<<<grid_size, BLOCK_SIZE>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                total_vecs,
                (int)total_elements
            );

            return output;
        }}
        """

        self.pad_op = load_inline(
            name="circularpad3d_static_opt",
            cpp_sources=cpp_header,
            cuda_sources=cuda_source,
            functions=["circularpad3d_cuda"],
            verbose=False,
            extra_cuda_cflags=["-O3", "--use_fast_math"]
        )

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