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.k = 5
        self.sigma_c = 10.0
        self.sigma_s = 5.0
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor bilateral_cuda(torch::Tensor input, int k_size, float sigma_color, float sigma_space);
        """
        
        cuda_source = """
        #include <cuda_runtime.h>
        #include <cmath>

        #define TILE_W 16
        #define TILE_H 16
        #define RADIUS 2  // (5-1)/2
        
        #define SMEM_W (TILE_W + 2 * RADIUS) // 20
        #define SMEM_H (TILE_H + 2 * RADIUS) // 20

        __global__ void bilateral_kernel(
            const float* __restrict__ input,
            float* __restrict__ output,
            int batch, int channels, int height, int width,
            double neg_inv_two_sigma_s_sq, // -1 / (2 * sigma_s^2)
            double neg_inv_two_sigma_c_sq  // -1 / (2 * sigma_c^2)
        ) {
            int tx = threadIdx.x;
            int ty = threadIdx.y;
            
            int ow = blockIdx.x * TILE_W + tx;
            int oh = blockIdx.y * TILE_H + ty;
            
            int bc = blockIdx.z;
            int c = bc % channels;
            int b = bc / channels;

            // Shared Memory
            __shared__ float smem[SMEM_H][SMEM_W];

         
            int tid = ty * TILE_W + tx; // 0..255
            int num_smem_pixels = SMEM_H * SMEM_W; // 400
            
            // Top-Left of the input window
            int in_h_base = blockIdx.y * TILE_H - RADIUS;
            int in_w_base = blockIdx.x * TILE_W - RADIUS;
            
            long long plane_offset = (long long)b * (channels * height * width) + c * (height * width);

            for (int i = tid; i < num_smem_pixels; i += 256) {
                int sh = i / SMEM_W;
                int sw = i % SMEM_W;
                
                int gh = in_h_base + sh;
                int gw = in_w_base + sw;
                
                // Clamp Padding logic (Replicate Border)
                int gh_clamped = min(max(gh, 0), height - 1);
                int gw_clamped = min(max(gw, 0), width - 1);
                
                smem[sh][sw] = input[plane_offset + gh_clamped * width + gw_clamped];
            }

            __syncthreads();

            if (ow < width && oh < height) {
                // Center pixel in smem
                // output (tx, ty) corresponds to smem (ty+R, tx+R)
                int center_sh = ty + RADIUS;
                int center_sw = tx + RADIUS;
                
                float center_val_f = smem[center_sh][center_sw];
                double center_val = (double)center_val_f;
                
                double sum_val = 0.0;
                double sum_weight = 0.0;

                // Loop over 5x5 window
                for (int ky = -RADIUS; ky <= RADIUS; ++ky) {
                    for (int kx = -RADIUS; kx <= RADIUS; ++kx) {
                        // Neighbor in smem
                        float neighbor_val_f = smem[center_sh + ky][center_sw + kx];
                        double neighbor_val = (double)neighbor_val_f;
                        
                        // 1. Spatial Weight
                        double dist_sq = (double)(ky*ky + kx*kx);
                        double w_space = exp(dist_sq * neg_inv_two_sigma_s_sq);
                        
                        // 2. Color Weight
                        double diff = neighbor_val - center_val;
                        double w_color = exp(diff * diff * neg_inv_two_sigma_c_sq);
                        
                        // Total Weight
                        double w = w_space * w_color;
                        
                        sum_val += neighbor_val * w;
                        sum_weight += w;
                    }
                }
                
                // 4. Write Output
                long long out_idx = plane_offset + oh * width + ow;
                output[out_idx] = (float)(sum_val / sum_weight);
            }
        }

        torch::Tensor bilateral_cuda(torch::Tensor input, int k_size, float sigma_color, float sigma_space) {
            int batch = input.size(0);
            int channels = input.size(1);
            int height = input.size(2);
            int width = input.size(3);
            
            auto output = torch::empty_like(input);
            
            // Pre-calculate constants (double)
            double neg_inv_two_sigma_s_sq = -1.0 / (2.0 * sigma_space * sigma_space);
            double neg_inv_two_sigma_c_sq = -1.0 / (2.0 * sigma_color * sigma_color);
            
            dim3 block(TILE_W, TILE_H);
            dim3 grid(
                (width + TILE_W - 1) / TILE_W,
                (height + TILE_H - 1) / TILE_H,
                batch * channels
            );

            bilateral_kernel<<<grid, block>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                batch, channels, height, width,
                neg_inv_two_sigma_s_sq,
                neg_inv_two_sigma_c_sq
            );
            
            return output;
        }
        """

        self.op = load_inline(
            name="bilateral_tiled_double_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["bilateral_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.bilateral_cuda(x, self.k, self.sigma_c, self.sigma_s)