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


from constantpad1d_torch import BATCH_SIZE, CHANNELS, W_IN, W_OUT, PADDING, PAD_VAL

PAD_L, PAD_R = PADDING

BLOCK_SIZE = 256
VEC_SIZE = 4 

class ModelNew(nn.Module):
    

    def __init__(self, padding, value):
        super().__init__()
        self.pad_l = padding[0] # 3
        self.pad_r = padding[1] # 1
        self.value = value      # 3.5
        
        self.w_in = W_IN
        self.w_out = W_OUT
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        
        cpp_header = """
        #include <torch/extension.h>

        torch::Tensor constantpad1d_cuda(
            torch::Tensor input,
            int pad_l, int pad_r, float value
        );
        """

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

        #define BLOCK_SIZE {BLOCK_SIZE}
        #define VEC_SIZE 4

        __global__ void constantpad1d_kernel(
            const float* __restrict__ input_data,
            float* __restrict__ output_data,
            int total_vecs,
            int total_elements,
            int W_in,
            int W_out,
            int pad_l,
            float value
        ) {{
            int idx = blockIdx.x * blockDim.x + threadIdx.x;
            int stride = gridDim.x * blockDim.x;

            float4* out_ptr = (float4*)output_data;

            for (int i = idx; i < total_vecs; i += stride) {{
                int start_idx = i * VEC_SIZE;
                
                int row = start_idx / W_out;
                int col = start_idx % W_out;
                
                int input_row_offset = row * W_in;
                float4 val_vec;

                #pragma unroll
                for (int k = 0; k < VEC_SIZE; k++) {{
                    int cur_w = col + k;
    
                    int current_input_offset = input_row_offset;
                    int effective_w = cur_w;
                    
                    if (effective_w >= W_out) {{
                        effective_w -= W_out;
                        current_input_offset += W_in;
                    }}

                    float val;

                    if (effective_w < pad_l) {{
                        val = value;
                    }} else if (effective_w >= pad_l + W_in) {{
                        val = value; 
                    }} else {{

                        val = input_data[current_input_offset + (effective_w - pad_l)];
                    }}

                    if (k == 0) val_vec.x = val;
                    else if (k == 1) val_vec.y = val;
                    else if (k == 2) val_vec.z = val;
                    else val_vec.w = val;
                }}
                out_ptr[i] = val_vec;
            }}


            int tail_start = total_vecs * VEC_SIZE;
            for (int k = tail_start + idx; k < total_elements; k += stride) {{
                 int row = k / W_out;
                 int w = k % W_out;
                 float val;
                 if (w < pad_l || w >= pad_l + W_in) val = value;
                 else val = input_data[row * W_in + (w - pad_l)];
                 output_data[k] = val;
            }}
        }}

        torch::Tensor constantpad1d_cuda(
            torch::Tensor input,
            int pad_l, int pad_r, float value
        ) {{
            input = input.contiguous();
            int N = input.size(0);
            int C = input.size(1);
            int W_in_sz = input.size(2);
            int W_out_sz = W_in_sz + pad_l + pad_r;
            
            auto output = torch::empty({{N, C, W_out_sz}}, input.options());
            int total_elements = N * C * W_out_sz;

            int total_vecs = total_elements / VEC_SIZE;
            int grid_size = (total_vecs + BLOCK_SIZE - 1) / BLOCK_SIZE;
            if (grid_size > 65535) grid_size = 65535;
            if (grid_size == 0) grid_size = 1;

            constantpad1d_kernel<<<grid_size, BLOCK_SIZE>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                total_vecs,
                total_elements,
                W_in_sz,
                W_out_sz,
                pad_l,
                value
            );
            return output;
        }}
        """

        self.pad_op = load_inline(
            name="constantpad1d_custom",
            cpp_sources=cpp_header,
            cuda_sources=cuda_source,
            functions=["constantpad1d_cuda"],
            verbose=False
        )

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