import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, applies Swish activation, 
    group normalization, and then HardSwish activation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, groups, eps, bias=True):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=bias)
        self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps)

    def forward(self, x):
        x = self.conv_transpose(x)
        x = torch.sigmoid(x) * x  # Swish activation
        x = self.group_norm(x)
        x = torch.nn.functional.hardswish(x)  # HardSwish activation
        return x

batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
groups = 4
eps = 1e-5

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, groups, eps]

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 bias,
                           torch::Tensor gn_w, torch::Tensor gn_b,
                           int64_t stride, int64_t padding, int64_t groups, double eps);
"""

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

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

__global__ void convtranspose3d_swish_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    float* __restrict__ y,
    int N, int Cin, int Din, int Hin, int Win,
    int Cout, int KD, int KH, int KW,
    int OD, int OH, int OW,
    int stride, int pad, int has_bias)
{
    int s = blockIdx.x * blockDim.x + threadIdx.x;
    int S = OD * OH * OW;
    if (s >= S) return;

    int nc = blockIdx.y;
    int co = nc % Cout;
    int n  = nc / Cout;

    int tmp = s;
    int ow = tmp % OW; tmp /= OW;
    int oh = tmp % OH;
    int od = tmp / OH;

    float acc = has_bias ? b[co] : 0.0f;

    #pragma unroll
    for (int ci = 0; ci < 3; ++ci) {
        if (ci >= Cin) break;
        for (int kd = 0; kd < KD; ++kd) {
            int id_num = od + pad - kd;
            if (id_num < 0 || (id_num % stride) != 0) continue;
            int id = id_num / stride;
            if ((unsigned)id >= (unsigned)Din) continue;
            for (int kh = 0; kh < KH; ++kh) {
                int ih_num = oh + pad - kh;
                if (ih_num < 0 || (ih_num % stride) != 0) continue;
                int ih = ih_num / stride;
                if ((unsigned)ih >= (unsigned)Hin) continue;
                for (int kw = 0; kw < KW; ++kw) {
                    int iw_num = ow + pad - kw;
                    if (iw_num < 0 || (iw_num % stride) != 0) continue;
                    int iw = iw_num / stride;
                    if ((unsigned)iw >= (unsigned)Win) continue;
                    float xv = x[(((n * Cin + ci) * Din + id) * Hin + ih) * Win + iw];
                    float wv = w[(((ci * Cout + co) * KD + kd) * KH + kh) * KW + kw];
                    acc = fmaf(xv, wv, acc);
                }
            }
        }
    }

    float sig = 1.0f / (1.0f + expf(-acc));
    y[((n * Cout + co) * S) + s] = acc * sig;
}

__global__ void group_reduce_kernel(
    const float* __restrict__ y,
    float* __restrict__ mean,
    float* __restrict__ invstd,
    int N, int C, int G, int S, float eps)
{
    extern __shared__ float sh[];
    float* shsum = sh;
    float* shsq  = sh + blockDim.x;

    int ng = blockIdx.x;
    int g = ng % G;
    int n = ng / G;
    int cpg = C / G;
    int L = cpg * S;
    int base = (n * C + g * cpg) * S;

    float sum = 0.0f;
    float sq = 0.0f;
    for (int i = threadIdx.x; i < L; i += blockDim.x) {
        float v = y[base + i];
        sum += v;
        sq = fmaf(v, v, sq);
    }
    shsum[threadIdx.x] = sum;
    shsq[threadIdx.x] = sq;
    __syncthreads();

    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (threadIdx.x < off) {
            shsum[threadIdx.x] += shsum[threadIdx.x + off];
            shsq[threadIdx.x]  += shsq[threadIdx.x + off];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        float m = shsum[0] / (float)L;
        float var = shsq[0] / (float)L - m * m;
        if (var < 0.0f) var = 0.0f;
        mean[ng] = m;
        invstd[ng] = rsqrtf(var + eps);
    }
}

__global__ void groupnorm_hardswish_kernel(
    float* __restrict__ y,
    const float* __restrict__ mean,
    const float* __restrict__ invstd,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    int N, int C, int G, int S)
{
    int s = blockIdx.x * blockDim.x + threadIdx.x;
    if (s >= S) return;
    int nc = blockIdx.y;
    int c = nc % C;
    int n = nc / C;
    int g = c / (C / G);

    int idx = (n * C + c) * S + s;
    float v = y[idx];
    float z = (v - mean[n * G + g]) * invstd[n * G + g];
    z = z * gamma[c] + beta[c];

    float t = z + 3.0f;
    float out;
    if (t <= 0.0f) out = 0.0f;
    else if (t >= 6.0f) out = z;
    else out = z * t * (1.0f / 6.0f);
    y[idx] = out;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                           torch::Tensor gn_w, torch::Tensor gn_b,
                           int64_t stride64, int64_t padding64, int64_t groups64, double eps_d) {
    CHECK_CUDA(x); CHECK_CUDA(w); CHECK_CUDA(gn_w); CHECK_CUDA(gn_b);
    CHECK_FLOAT(x); CHECK_FLOAT(w); CHECK_FLOAT(gn_w); CHECK_FLOAT(gn_b);
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(w.is_contiguous(), "w must be contiguous");

    int N = (int)x.size(0);
    int Cin = (int)x.size(1);
    int Din = (int)x.size(2);
    int Hin = (int)x.size(3);
    int Win = (int)x.size(4);
    int Cout = (int)w.size(1);
    int KD = (int)w.size(2);
    int KH = (int)w.size(3);
    int KW = (int)w.size(4);
    int stride = (int)stride64;
    int pad = (int)padding64;
    int G = (int)groups64;
    int OD = (Din - 1) * stride - 2 * pad + KD;
    int OH = (Hin - 1) * stride - 2 * pad + KH;
    int OW = (Win - 1) * stride - 2 * pad + KW;
    int S = OD * OH * OW;

    auto y = torch::empty({N, Cout, OD, OH, OW}, x.options());
    auto stats_opts = x.options();
    auto mean = torch::empty({N * G}, stats_opts);
    auto invstd = torch::empty({N * G}, stats_opts);

    const int threads = 256;
    dim3 grid_conv((S + threads - 1) / threads, N * Cout);
    int has_bias = (bias.defined() && bias.numel() > 0) ? 1 : 0;
    const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;
    convtranspose3d_swish_kernel<<<grid_conv, threads>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bptr, y.data_ptr<float>(),
        N, Cin, Din, Hin, Win, Cout, KD, KH, KW, OD, OH, OW, stride, pad, has_bias);

    int red_threads = 256;
    group_reduce_kernel<<<N * G, red_threads, red_threads * 2 * sizeof(float)>>>(
        y.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
        N, Cout, G, S, (float)eps_d);

    dim3 grid_apply((S + threads - 1) / threads, N * Cout);
    groupnorm_hardswish_kernel<<<grid_apply, threads>>>(
        y.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
        gn_w.data_ptr<float>(), gn_b.data_ptr<float>(), N, Cout, G, S);

    return y;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, groups, eps, bias=True):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=bias)
        self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps)
        self._groups = groups
        self._eps = eps
        if isinstance(stride, tuple):
            self._stride = int(stride[0])
        else:
            self._stride = int(stride)
        if isinstance(padding, tuple):
            self._padding = int(padding[0])
        else:
            self._padding = int(padding)
        self.register_buffer("_empty_bias", torch.empty(0), persistent=False)

    def forward(self, x):
        if not x.is_contiguous():
            x = x.contiguous()
        b = self.conv_transpose.bias
        if b is None:
            b = self._empty_bias
        return _ext.fused_forward(x, self.conv_transpose.weight, b,
                                  self.group_norm.weight, self.group_norm.bias,
                                  self._stride, self._padding, self._groups, float(self._eps))