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

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

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, value):
        super().__init__()
        self.value = value
        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 {PAD_L}
        #define PAD_T {PAD_T}
        #define PAD_F {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 constantpad3d_cuda(torch::Tensor input, float value);
        """

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

        {macros}

        __global__ void __launch_bounds__(BLOCK_SIZE) constantpad3d_kernel_static(
            const float* __restrict__ input_data,
            float* __restrict__ output_data,
            int total_vecs,
            int total_elements,
            float value
        ) {{
            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) {{
                
                // Compile-time unroll
                #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 = start_linear_idx % W_OUT;
                    int tmp   = start_linear_idx / 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 input_base = nc_idx * DHW_IN;

                    float4 val_vec;

                    #pragma unroll
                    for (int v = 0; v < VEC_SIZE; ++v) {{
                        int cur_w = w_out + v;
                        int cur_h = h_out;
                        int cur_d = d_out;
                        int cur_input_base = input_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;
                                    cur_input_base += DHW_IN; 
                                }}
                            }}
                        }}

                        float val;

   
                        // 1. Check Front/Back
                        if ((unsigned int)(cur_d - PAD_F) >= (unsigned int)D_IN) {{
                            val = value;
                        }}
                        // 2. Check Top/Bottom
                        else if ((unsigned int)(cur_h - PAD_T) >= (unsigned int)H_IN) {{
                            val = value;
                        }}
                        // 3. Check Left/Right
                        else if ((unsigned int)(cur_w - PAD_L) >= (unsigned int)W_IN) {{
                            val = value;
                        }}
                        // 4. Copy from Input
                        else {{
                 
                            int d_off = cur_d - PAD_F;
                            int h_off = cur_h - PAD_T;
                            int w_off = cur_w - PAD_L;
                            
                            int read_idx = cur_input_base + (d_off * HW_IN) + (h_off * W_IN) + w_off;
                            val = input_data[read_idx];
                        }}

                        if (v == 0) val_vec.x = val;
                        else if (v == 1) val_vec.y = val;
                        else if (v == 2) val_vec.z = val;
                        else val_vec.w = val;
                    }}

                    // Vectorized Store
                    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;

                 float val;
                 if ((unsigned int)(d_out - PAD_F) >= (unsigned int)D_IN) val = value;
                 else if ((unsigned int)(h_out - PAD_T) >= (unsigned int)H_IN) val = value;
                 else if ((unsigned int)(w_out - PAD_L) >= (unsigned int)W_IN) val = value;
                 else {{
                     int read_idx = nc_idx * DHW_IN 
                                  + (d_out - PAD_F) * HW_IN 
                                  + (h_out - PAD_T) * W_IN 
                                  + (w_out - PAD_L);
                     val = input_data[read_idx];
                 }}
                 output_data[k] = val;
            }}
        }}

        torch::Tensor constantpad3d_cuda(torch::Tensor input, float value) {{
      
            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;

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

            return output;
        }}
        """

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

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