import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    """
    Model that performs:
    1. Conv3D
    2. HardSwish activation
    3. GroupNorm  
    4. Mean pooling across spatial dimensions
    """
    def __init__(self, in_channels, out_channels, kernel_size, num_groups=4, bias=True):
        super(Model, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, bias=bias)
        self.group_norm = nn.GroupNorm(num_groups, out_channels)

    def forward(self, x):
        x = self.conv(x)                             # (B, C, D, H, W)
        x = F.hardswish(x)                           # Nonlinear activation
        x = self.group_norm(x)                       # Normalization over channels
        x = torch.mean(x, dim=[2, 3, 4])             # Mean over spatial dims → (B, C)
        return x

# === Test config ===
batch_size = 1024
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 4

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]

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

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

torch::Tensor fused_conv_hswish_gn_pool(torch::Tensor x,
                                        torch::Tensor conv_w,
                                        torch::Tensor conv_b,
                                        torch::Tensor gn_w,
                                        torch::Tensor gn_b,
                                        int64_t groups);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

__device__ __forceinline__ float hswish_f(float x) {
    float t = x + 3.0f;
    t = t < 0.0f ? 0.0f : (t > 6.0f ? 6.0f : t);
    return x * t * (1.0f / 6.0f);
}

// Specialized for the benchmark Conv3d kernel_size = 4 and group size = 4 channels.
__global__ __launch_bounds__(256, 2)
void fused_k4_g4_kernel(const float* __restrict__ x,
                        const float* __restrict__ w,
                        const float* __restrict__ bias,
                        const float* __restrict__ gn_w,
                        const float* __restrict__ gn_b,
                        float* __restrict__ out,
                        int B, int Cin, int D, int H, int W,
                        int Cout, int groups, int has_bias) {
    const int tid = threadIdx.x;
    const int bg = blockIdx.x;
    const int b = bg / groups;
    const int g = bg - b * groups;

    const int Do = D - 3;
    const int Ho = H - 3;
    const int Wo = W - 3;
    const int S = Do * Ho * Wo;
    const int cbase = g * 4;

    float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
    float ss = 0.0f;

    for (int p = tid; p < S; p += blockDim.x) {
        int tmp = p;
        const int ow = tmp % Wo; tmp /= Wo;
        const int oh = tmp % Ho;
        const int od = tmp / Ho;

        float a0 = has_bias ? bias[cbase + 0] : 0.0f;
        float a1 = has_bias ? bias[cbase + 1] : 0.0f;
        float a2 = has_bias ? bias[cbase + 2] : 0.0f;
        float a3 = has_bias ? bias[cbase + 3] : 0.0f;

        #pragma unroll
        for (int ci = 0; ci < 3; ++ci) {
            const float* __restrict__ xp = x + (((b * Cin + ci) * D + od) * H + oh) * W + ow;
            const float* __restrict__ w0 = w + (((cbase + 0) * Cin + ci) * 4) * 4 * 4;
            const float* __restrict__ w1 = w + (((cbase + 1) * Cin + ci) * 4) * 4 * 4;
            const float* __restrict__ w2 = w + (((cbase + 2) * Cin + ci) * 4) * 4 * 4;
            const float* __restrict__ w3 = w + (((cbase + 3) * Cin + ci) * 4) * 4 * 4;
            #pragma unroll
            for (int kd = 0; kd < 4; ++kd) {
                #pragma unroll
                for (int kh = 0; kh < 4; ++kh) {
                    const float* __restrict__ xrow = xp + kd * H * W + kh * W;
                    const int koff = (kd * 4 + kh) * 4;
                    #pragma unroll
                    for (int kw = 0; kw < 4; ++kw) {
                        float xv = xrow[kw];
                        a0 = fmaf(xv, w0[koff + kw], a0);
                        a1 = fmaf(xv, w1[koff + kw], a1);
                        a2 = fmaf(xv, w2[koff + kw], a2);
                        a3 = fmaf(xv, w3[koff + kw], a3);
                    }
                }
            }
        }

        float y0 = hswish_f(a0);
        float y1 = hswish_f(a1);
        float y2 = hswish_f(a2);
        float y3 = hswish_f(a3);
        s0 += y0; s1 += y1; s2 += y2; s3 += y3;
        ss = fmaf(y0, y0, ss);
        ss = fmaf(y1, y1, ss);
        ss = fmaf(y2, y2, ss);
        ss = fmaf(y3, y3, ss);
    }

    extern __shared__ float sh[];
    sh[tid] = s0;
    sh[tid + blockDim.x] = s1;
    sh[tid + 2 * blockDim.x] = s2;
    sh[tid + 3 * blockDim.x] = s3;
    sh[tid + 4 * blockDim.x] = ss;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            sh[tid] += sh[tid + stride];
            sh[tid + blockDim.x] += sh[tid + blockDim.x + stride];
            sh[tid + 2 * blockDim.x] += sh[tid + 2 * blockDim.x + stride];
            sh[tid + 3 * blockDim.x] += sh[tid + 3 * blockDim.x + stride];
            sh[tid + 4 * blockDim.x] += sh[tid + 4 * blockDim.x + stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        const float rs = 1.0f / (float)S;
        const float invN = 1.0f / (4.0f * (float)S);
        float cs0 = sh[0];
        float cs1 = sh[blockDim.x];
        float cs2 = sh[2 * blockDim.x];
        float cs3 = sh[3 * blockDim.x];
        float gsum = cs0 + cs1 + cs2 + cs3;
        float mean = gsum * invN;
        float var = sh[4 * blockDim.x] * invN - mean * mean;
        var = var < 0.0f ? 0.0f : var;
        float rstd = rsqrtf(var + 1.0e-5f);

        out[b * Cout + cbase + 0] = gn_w[cbase + 0] * ((cs0 * rs - mean) * rstd) + gn_b[cbase + 0];
        out[b * Cout + cbase + 1] = gn_w[cbase + 1] * ((cs1 * rs - mean) * rstd) + gn_b[cbase + 1];
        out[b * Cout + cbase + 2] = gn_w[cbase + 2] * ((cs2 * rs - mean) * rstd) + gn_b[cbase + 2];
        out[b * Cout + cbase + 3] = gn_w[cbase + 3] * ((cs3 * rs - mean) * rstd) + gn_b[cbase + 3];
    }
}

torch::Tensor fused_conv_hswish_gn_pool(torch::Tensor x,
                                        torch::Tensor conv_w,
                                        torch::Tensor conv_b,
                                        torch::Tensor gn_w,
                                        torch::Tensor gn_b,
                                        int64_t groups64) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    x = x.contiguous();
    conv_w = conv_w.contiguous();
    gn_w = gn_w.contiguous();
    gn_b = gn_b.contiguous();
    if (conv_b.numel() > 0) conv_b = conv_b.contiguous();

    const int B = (int)x.size(0);
    const int Cin = (int)x.size(1);
    const int D = (int)x.size(2);
    const int H = (int)x.size(3);
    const int W = (int)x.size(4);
    const int Cout = (int)conv_w.size(0);
    const int groups = (int)groups64;
    TORCH_CHECK(Cout / groups == 4, "specialized kernel requires 4 channels per group");
    TORCH_CHECK(conv_w.size(2) == 4 && conv_w.size(3) == 4 && conv_w.size(4) == 4, "specialized kernel requires kernel_size=4");
    TORCH_CHECK(Cin == 3, "specialized kernel requires in_channels=3");

    auto out = torch::empty({B, Cout}, x.options());
    const int threads = 256;
    const dim3 grid(B * groups);
    const size_t smem = 5 * threads * sizeof(float);
    const int has_bias = conv_b.numel() > 0 ? 1 : 0;
    const float* bptr = has_bias ? conv_b.data_ptr<float>() : nullptr;

    fused_k4_g4_kernel<<<grid, threads, smem, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), conv_w.data_ptr<float>(), bptr,
        gn_w.data_ptr<float>(), gn_b.data_ptr<float>(), out.data_ptr<float>(),
        B, Cin, D, H, W, Cout, groups, has_bias);
    return out;
}
"""

_ext = load_inline(
    name="fused_conv3d_hswish_gn_pool_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_conv_hswish_gn_pool"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, num_groups=4, bias=True):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, bias=bias)
        self.group_norm = nn.GroupNorm(num_groups, out_channels)
        self.num_groups = num_groups

    def forward(self, x):
        if self.conv.bias is None:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
        else:
            b = self.conv.bias
        return _ext.fused_conv_hswish_gn_pool(x, self.conv.weight, b,
                                               self.group_norm.weight, self.group_norm.bias,
                                               self.num_groups)