import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs convolution, group normalization, scaling, max pooling, and clamping.
    """
    def __init__(self, in_channels, out_channels, kernel_size, num_groups, scale_shape, maxpool_kernel_size, clamp_min, clamp_max):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.group_norm = nn.GroupNorm(num_groups, out_channels)
        self.scale = nn.Parameter(torch.ones(scale_shape))
        self.maxpool = nn.MaxPool2d(kernel_size=maxpool_kernel_size)
        self.clamp_min = clamp_min
        self.clamp_max = clamp_max

    def forward(self, x):
        """
        Args:
            x: Input tensor of shape (batch_size, in_channels, height, width).
        Returns:
            Output tensor of shape (batch_size, out_channels, height', width').
        """
        x = self.conv(x)
        x = self.group_norm(x)
        x = x * self.scale
        x = self.maxpool(x)
        x = torch.clamp(x, self.clamp_min, self.clamp_max)
        return x

batch_size = 128
in_channels = 8
out_channels = 64
height, width = 128, 128 
kernel_size = 3
num_groups = 16
scale_shape = (out_channels, 1, 1)
maxpool_kernel_size = 4
clamp_min = 0.0
clamp_max = 1.0

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, num_groups, scale_shape, maxpool_kernel_size, clamp_min, clamp_max]

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

_cpp_src = r'''
#include <torch/extension.h>
#include <vector>

void gn_stats_cuda(torch::Tensor x, torch::Tensor stats, int64_t groups);
void gn_pool_out_cuda(torch::Tensor x, torch::Tensor stats, torch::Tensor gamma, torch::Tensor beta,
                      torch::Tensor scale, torch::Tensor out, int64_t groups, int64_t pool,
                      double clamp_min, double clamp_max);

torch::Tensor gn_pool_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta,
                              torch::Tensor scale, int64_t groups, int64_t pool,
                              double clamp_min, double clamp_max) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(gamma.is_cuda() && beta.is_cuda() && scale.is_cuda(), "params must be CUDA");
    const int64_t N = x.size(0);
    const int64_t C = x.size(1);
    const int64_t H = x.size(2);
    const int64_t W = x.size(3);
    const int64_t OH = (H - pool) / pool + 1;
    const int64_t OW = (W - pool) / pool + 1;
    auto stats = torch::empty({N, groups, 2}, x.options());
    auto out = torch::empty({N, C, OH, OW}, x.options());
    gn_stats_cuda(x, stats, groups);
    gn_pool_out_cuda(x, stats, gamma.contiguous(), beta.contiguous(), scale.contiguous(), out,
                     groups, pool, clamp_min, clamp_max);
    return out;
}
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>

__global__ void gn_stats_kernel(const float* __restrict__ x, float* __restrict__ stats,
                                int N, int C, int H, int W, int G) {
    const int bg = blockIdx.x;
    const int n = bg / G;
    const int g = bg - n * G;
    const int tid = threadIdx.x;
    const int HW = H * W;
    const int CG = C / G;
    const int total = CG * HW;
    float sum = 0.0f;
    float sum2 = 0.0f;
    const int c0 = g * CG;
    const int baseN = n * C * HW;

    for (int i = tid; i < total; i += blockDim.x) {
        int cg = i / HW;
        int hw = i - cg * HW;
        float v = x[baseN + (c0 + cg) * HW + hw];
        sum += v;
        sum2 += v * v;
    }

    __shared__ float ssum[256];
    __shared__ float ssq[256];
    ssum[tid] = sum;
    ssq[tid] = sum2;
    __syncthreads();

    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) {
            ssum[tid] += ssum[tid + s];
            ssq[tid] += ssq[tid + s];
        }
        __syncthreads();
    }
    if (tid == 0) {
        float inv_count = 1.0f / (float)total;
        float mean = ssum[0] * inv_count;
        float var = ssq[0] * inv_count - mean * mean;
        var = var < 0.0f ? 0.0f : var;
        float invstd = rsqrtf(var + 1.0e-5f);
        int off = bg * 2;
        stats[off] = mean;
        stats[off + 1] = invstd;
    }
}

__global__ void gn_pool_out_kernel(const float* __restrict__ x, const float* __restrict__ stats,
                                   const float* __restrict__ gamma, const float* __restrict__ beta,
                                   const float* __restrict__ scale, float* __restrict__ out,
                                   int N, int C, int H, int W, int G, int pool,
                                   int OH, int OW, float clamp_min, float clamp_max) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    const int64_t total = (int64_t)N * C * OH * OW;
    if (idx >= total) return;

    int ow = idx % OW;
    int64_t t = idx / OW;
    int oh = t % OH;
    t /= OH;
    int c = t % C;
    int n = t / C;

    int CG = C / G;
    int g = c / CG;
    int stat_off = (n * G + g) * 2;
    float mean = stats[stat_off];
    float invstd = stats[stat_off + 1];
    float gm = gamma[c];
    float bt = beta[c];
    float sc = scale[c];
    int HW = H * W;
    int base = (n * C + c) * HW + (oh * pool) * W + (ow * pool);

    float m = -FLT_MAX;
    #pragma unroll
    for (int ph = 0; ph < 4; ++ph) {
        #pragma unroll
        for (int pw = 0; pw < 4; ++pw) {
            float v = x[base + ph * W + pw];
            v = ((v - mean) * invstd * gm + bt) * sc;
            m = fmaxf(m, v);
        }
    }
    // Correct for non-4 pool sizes if ever used with same API.
    if (pool != 4) {
        m = -FLT_MAX;
        for (int ph = 0; ph < pool; ++ph) {
            for (int pw = 0; pw < pool; ++pw) {
                float v = x[base + ph * W + pw];
                v = ((v - mean) * invstd * gm + bt) * sc;
                m = fmaxf(m, v);
            }
        }
    }
    m = fminf(fmaxf(m, clamp_min), clamp_max);
    out[idx] = m;
}

void gn_stats_cuda(torch::Tensor x, torch::Tensor stats, int64_t groups) {
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const dim3 blocks(N * (int)groups);
    const dim3 threads(256);
    gn_stats_kernel<<<blocks, threads>>>(x.data_ptr<float>(), stats.data_ptr<float>(), N, C, H, W, (int)groups);
}

void gn_pool_out_cuda(torch::Tensor x, torch::Tensor stats, torch::Tensor gamma, torch::Tensor beta,
                      torch::Tensor scale, torch::Tensor out, int64_t groups, int64_t pool,
                      double clamp_min, double clamp_max) {
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int OH = (int)out.size(2);
    const int OW = (int)out.size(3);
    const int64_t total = (int64_t)N * C * OH * OW;
    const int threads = 256;
    const int blocks = (int)((total + threads - 1) / threads);
    gn_pool_out_kernel<<<blocks, threads>>>(x.data_ptr<float>(), stats.data_ptr<float>(),
        gamma.data_ptr<float>(), beta.data_ptr<float>(), scale.data_ptr<float>(), out.data_ptr<float>(),
        N, C, H, W, (int)groups, (int)pool, OH, OW, (float)clamp_min, (float)clamp_max);
}
'''

_ext = load_inline(
    name="fused_gn_pool_clamp_ext_v1",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["gn_pool_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, num_groups, scale_shape,
                 maxpool_kernel_size, clamp_min, clamp_max):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.group_norm = nn.GroupNorm(num_groups, out_channels)
        self.scale = nn.Parameter(torch.ones(scale_shape))
        self.maxpool_kernel_size = maxpool_kernel_size
        self.clamp_min = float(clamp_min)
        self.clamp_max = float(clamp_max)

    def forward(self, x):
        y = F.conv2d(x, self.conv.weight, self.conv.bias, stride=1, padding=0, dilation=1, groups=1)
        y = y.contiguous()
        return _ext.gn_pool_forward(y, self.group_norm.weight, self.group_norm.bias,
                                    self.scale, self.group_norm.num_groups,
                                    int(self.maxpool_kernel_size), self.clamp_min, self.clamp_max)