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

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.kernel_size = 3
        self.stride = 1
        self.padding = 1
        self.height = 64
        self.width = 64
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor fold_cuda(torch::Tensor input, int output_h, int output_w, int k_size, int stride, int padding);
        """
        
        cuda_source = """
        #include <cuda_runtime.h>

        #define BLOCK_SIZE 256

        __global__ void fold_gather_f4_kernel(
            const float* __restrict__ input,
            float* __restrict__ output,
            int n_vec,   
            int batch,
            int channels,
            int out_h,
            int out_w,
            int k_size,
            int stride,
            int padding,
            int L_in     
        ) {
            int idx = blockIdx.x * blockDim.x + threadIdx.x;
            if (idx >= n_vec) return;

            int w_vec_dim = out_w / 4;
            
            int tmp = idx;
            int w_vec = tmp % w_vec_dim; tmp /= w_vec_dim;
            int h = tmp % out_h; tmp /= out_h;
            int c = tmp % channels;
            int b = tmp / channels;

            int w_start = w_vec * 4;

            float sum0 = 0.0f, sum1 = 0.0f, sum2 = 0.0f, sum3 = 0.0f;

            long long input_batch_offset = (long long)b * (channels * k_size * k_size * L_in);

            for (int kh = 0; kh < k_size; ++kh) {
                int dist_h = h + padding - kh;
                
                if (dist_h % stride == 0) {
                    int win_h = dist_h / stride;
                    
                    int grid_w = (out_w + 2 * padding - k_size) / stride + 1;
                    int grid_h = (out_h + 2 * padding - k_size) / stride + 1;

                    if (win_h >= 0 && win_h < grid_h) {
                        
                       
                        for (int kw = 0; kw < k_size; ++kw) {
                            
                            // Input Channel Offset: c * K*K + kh*K + kw
                            int c_in_idx = c * k_size * k_size + kh * k_size + kw;
                            long long input_channel_ptr = input_batch_offset + (long long)c_in_idx * L_in;

                            int dist_w0 = w_start + 0 + padding - kw;
                            int dist_w1 = w_start + 1 + padding - kw;
                            int dist_w2 = w_start + 2 + padding - kw;
                            int dist_w3 = w_start + 3 + padding - kw;

                            // Pixel 0
                            if (dist_w0 % stride == 0) {
                                int win_w = dist_w0 / stride;
                                if (win_w >= 0 && win_w < grid_w) {
                                    int l_idx = win_h * grid_w + win_w;
                                    sum0 += input[input_channel_ptr + l_idx];
                                }
                            }
                            // Pixel 1
                            if (dist_w1 % stride == 0) {
                                int win_w = dist_w1 / stride;
                                if (win_w >= 0 && win_w < grid_w) {
                                    int l_idx = win_h * grid_w + win_w;
                                    sum1 += input[input_channel_ptr + l_idx];
                                }
                            }
                            // Pixel 2
                            if (dist_w2 % stride == 0) {
                                int win_w = dist_w2 / stride;
                                if (win_w >= 0 && win_w < grid_w) {
                                    int l_idx = win_h * grid_w + win_w;
                                    sum2 += input[input_channel_ptr + l_idx];
                                }
                            }
                            // Pixel 3
                            if (dist_w3 % stride == 0) {
                                int win_w = dist_w3 / stride;
                                if (win_w >= 0 && win_w < grid_w) {
                                    int l_idx = win_h * grid_w + win_w;
                                    sum3 += input[input_channel_ptr + l_idx];
                                }
                            }
                        }
                    }
                }
            }

            long long out_idx = (long long)b * (channels * out_h * out_w) + c * (out_h * out_w) + h * out_w + w_start;
            
            output[out_idx + 0] = sum0;
            output[out_idx + 1] = sum1;
            output[out_idx + 2] = sum2;
            output[out_idx + 3] = sum3;
        }

        torch::Tensor fold_cuda(torch::Tensor input, int output_h, int output_w, int k_size, int stride, int padding) {
            int batch = input.size(0);
            int c_total = input.size(1); // C * K * K
            int L_in = input.size(2);
            
            int channels = c_total / (k_size * k_size);
            
            auto output = torch::empty({batch, channels, output_h, output_w}, input.options());
            
            if (output_w % 4 != 0) return output;

            long long total_elements = output.numel();
            int n_vec = total_elements / 4;
            
            const int block_size = 256;
            const int grid_size = (n_vec + block_size - 1) / block_size;

            fold_gather_f4_kernel<<<grid_size, block_size>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                n_vec,
                batch, channels, output_h, output_w,
                k_size, stride, padding, L_in
            );
            
            return output;
        }
        """

        self.op = load_inline(
            name="fold_gather_f4_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["fold_cuda"],
            extra_cuda_cflags=["-O3"],
            verbose=False
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_contiguous(): x = x.contiguous()
        return self.op.fold_cuda(x, self.height, self.width, self.kernel_size, self.stride, self.padding)