# localresponsenorm_cuda.py
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

from localresponsenorm_torch import N, C, H, W, SIZE, ALPHA, BETA, K

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        half_size = SIZE // 2
        
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor lrn_forward_cuda(torch::Tensor input, int size, int half_size, float alpha, float beta, float k);
        """

        cuda_source = """
        #include <cuda_runtime.h>
        #include <device_launch_parameters.h>

        #define BLOCK_SIZE 256

        __global__ void lrn_fused_kernel(
            const float* __restrict__ x,
            float* __restrict__ y,
            int N, int C, int H, int W,
            int size, int half_size, float alpha, float beta, float k
        ) {
            int idx = blockIdx.x * blockDim.x + threadIdx.x;
            if (idx >= N * C * H * W) return;

            int n = idx / (C * H * W);
            int rem = idx % (C * H * W);
            int c = rem / (H * W);
            int hw = rem % (H * W);

            int base_idx = n * C * H * W + hw;
            float sum_sq = 0.0f;
            
            int start_c = max(0, c - half_size);
            int end_c = min(C - 1, c + half_size);

            #pragma unroll 4
            for (int window_c = start_c; window_c <= end_c; ++window_c) {
                float val = x[base_idx + window_c * H * W];
                sum_sq = __fmaf_rn(val, val, sum_sq);
            }

            float alpha_over_n = alpha / (float)size;
            float scale = __fmaf_rn(alpha_over_n, sum_sq, k);
            

            float norm_factor = __powf(scale, beta);
            float input_val = x[idx];
            y[idx] = __fdividef(input_val, norm_factor);
        }

        torch::Tensor lrn_forward_cuda(torch::Tensor input, int size, int half_size, float alpha, float beta, float k) {
            const int N = input.size(0);
            const int C = input.size(1);
            const int H = input.size(2);
            const int W = input.size(3);
            auto output = torch::empty_like(input);

            const int n_elements = N * C * H * W;
            const int blocks = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;

            lrn_fused_kernel<<<blocks, BLOCK_SIZE>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                N, C, H, W,
                size, half_size, alpha, beta, k
            );
            
            return output;
        }
        """
        
        self.lrn_op = load_inline(
            name="lrn_correct_formula",
            cpp_sources=cpp_source, 
            cuda_sources=cuda_source,
            functions=["lrn_forward_cuda"],
            extra_cuda_cflags=[
                "-O3",
                "--use_fast_math",
                "--fmad=true"
            ],
            verbose=True
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        half_size = SIZE // 2
        return self.lrn_op.lrn_forward_cuda(x.contiguous(), SIZE, half_size, ALPHA, BETA, K)