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_forward(torch::Tensor x, torch::Tensor w, torch::Tensor conv_b, torch::Tensor gn_w, torch::Tensor gn_b, double eps, int64_t num_groups);
'''

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

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be CUDA")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")

__device__ __forceinline__ float hswish_f(float v) {
    if (v <= -3.0f) return 0.0f;
    if (v >= 3.0f) return v;
    return v * (v + 3.0f) * 0.16666666666666666f;
}

__global__ __launch_bounds__(256, 2) void group4_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 has_bias, float eps)
{
    constexpr int D=16, H=32, Wd=32, Do=13, Ho=29, Wo=29, S=10933;
    int bg = blockIdx.x;
    int b = bg >> 2;
    int g = bg & 3;
    int cbase = g << 2;
    int tid = threadIdx.x;

    extern __shared__ float smem[];
    float* sw = smem;
    float* rs = sw + 768;
    float* rq = rs + 1024;

    for (int i = tid; i < 768; i += 256) {
        int p = i >> 2;
        int oc = i & 3;
        sw[i] = w[(cbase + oc) * 192 + p];
    }
    __syncthreads();

    float s0=0.0f, s1=0.0f, s2=0.0f, s3=0.0f;
    float q0=0.0f, q1=0.0f, q2=0.0f, q3=0.0f;
    float b0=0.0f, b1=0.0f, b2=0.0f, b3=0.0f;
    if (has_bias) { b0=bias[cbase]; b1=bias[cbase+1]; b2=bias[cbase+2]; b3=bias[cbase+3]; }

    for (int idx = tid; idx < S; idx += 256) {
        int od = idx / 841;
        int rem = idx - od * 841;
        int oh = rem / 29;
        int ow = rem - oh * 29;
        float a0=b0, a1=b1, a2=b2, a3=b3;
        #pragma unroll
        for (int ci=0; ci<3; ++ci) {
            #pragma unroll
            for (int kd=0; kd<4; ++kd) {
                int id = od + kd;
                #pragma unroll
                for (int kh=0; kh<4; ++kh) {
                    int ih = oh + kh;
                    int xbase = (((b*3 + ci) * D + id) * H + ih) * Wd + ow;
                    int pbase = ci*64 + kd*16 + kh*4;
                    #pragma unroll
                    for (int kw=0; kw<4; ++kw) {
                        float xv = x[xbase + kw];
                        int pp = (pbase + kw) << 2;
                        a0 = fmaf(xv, sw[pp+0], a0);
                        a1 = fmaf(xv, sw[pp+1], a1);
                        a2 = fmaf(xv, sw[pp+2], a2);
                        a3 = fmaf(xv, sw[pp+3], a3);
                    }
                }
            }
        }
        float y0=hswish_f(a0), y1=hswish_f(a1), y2=hswish_f(a2), y3=hswish_f(a3);
        s0 += y0; s1 += y1; s2 += y2; s3 += y3;
        q0 = fmaf(y0,y0,q0); q1 = fmaf(y1,y1,q1); q2 = fmaf(y2,y2,q2); q3 = fmaf(y3,y3,q3);
    }

    rs[tid] = s0; rs[256+tid]=s1; rs[512+tid]=s2; rs[768+tid]=s3;
    rq[tid] = q0; rq[256+tid]=q1; rq[512+tid]=q2; rq[768+tid]=q3;
    __syncthreads();
    for (int stride=128; stride>0; stride >>= 1) {
        if (tid < stride) {
            rs[tid]+=rs[tid+stride]; rs[256+tid]+=rs[256+tid+stride]; rs[512+tid]+=rs[512+tid+stride]; rs[768+tid]+=rs[768+tid+stride];
            rq[tid]+=rq[tid+stride]; rq[256+tid]+=rq[256+tid+stride]; rq[512+tid]+=rq[512+tid+stride]; rq[768+tid]+=rq[768+tid+stride];
        }
        __syncthreads();
    }
    if (tid == 0) {
        float sum0=rs[0], sum1=rs[256], sum2=rs[512], sum3=rs[768];
        float sq0=rq[0], sq1=rq[256], sq2=rq[512], sq3=rq[768];
        constexpr float invS = 1.0f / 10933.0f;
        float mean = (sum0+sum1+sum2+sum3) * (0.25f * invS);
        float var = (sq0+sq1+sq2+sq3) * (0.25f * invS) - mean*mean;
        var = fmaxf(var, 0.0f);
        float inv = rsqrtf(var + eps);
        int ob = b * 16 + cbase;
        out[ob+0] = ((sum0*invS - mean) * inv) * gn_w[cbase+0] + gn_b[cbase+0];
        out[ob+1] = ((sum1*invS - mean) * inv) * gn_w[cbase+1] + gn_b[cbase+1];
        out[ob+2] = ((sum2*invS - mean) * inv) * gn_w[cbase+2] + gn_b[cbase+2];
        out[ob+3] = ((sum3*invS - mean) * inv) * gn_w[cbase+3] + gn_b[cbase+3];
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor conv_b, torch::Tensor gn_w, torch::Tensor gn_b, double eps, int64_t num_groups) {
    CHECK_CUDA(x); CHECK_CUDA(w); CHECK_CUDA(gn_w); CHECK_CUDA(gn_b);
    CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(w); CHECK_CONTIGUOUS(gn_w); CHECK_CONTIGUOUS(gn_b);
    int B=x.size(0);
    auto out=torch::empty({B, 16}, x.options());
    int has_bias = conv_b.numel() > 0;
    const float* bp = has_bias ? conv_b.data_ptr<float>() : nullptr;
    size_t shmem = (768 + 2048) * sizeof(float);
    group4_kernel<<<B*4, 256, shmem>>>(x.data_ptr<float>(), w.data_ptr<float>(), bp, gn_w.data_ptr<float>(), gn_b.data_ptr<float>(), out.data_ptr<float>(), has_bias, (float)eps);
    return out;
}
'''

_ext = load_inline(name='fused_conv3d_hswish_gn_mean_ext_v3', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['fused_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], extra_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)
    def forward(self, x):
        b = self.conv.bias
        if b is None:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.fused_forward(x.contiguous(), self.conv.weight.contiguous(), b.contiguous(), self.group_norm.weight.contiguous(), self.group_norm.bias.contiguous(), float(self.group_norm.eps), int(self.group_norm.num_groups))