import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a 3D transposed convolution operation with square input and square kernel,
    and supports padding, dilation, and stride.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the convolution kernel (square kernel, so only one value needed).
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        dilation (int, optional): Spacing between kernel elements. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(in_channels, out_channels, kernel_size=(kernel_size, kernel_size, kernel_size), stride=stride, padding=padding, dilation=dilation, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 3D transposed convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, depth, height, width).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, depth_out, height_out, width_out).
        """
        return self.conv_transpose3d(x)

# Test code
batch_size = 16
in_channels = 32
out_channels = 64
kernel_size = 3
depth = 16
height = 32
width = 32
stride = 2
padding = 1
dilation = 2

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

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

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

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

torch::Tensor deconv3d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                    int64_t stride, int64_t padding, int64_t dilation);

torch::Tensor deconv3d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                               int64_t stride, int64_t padding, int64_t dilation) {
    return deconv3d_forward_cuda(x, weight, bias, stride, padding, dilation);
}
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void init_bias_kernel(float* __restrict__ y, const float* __restrict__ bias,
                                 long total, int Cout, int outD, int outH, int outW) {
    long idx = blockIdx.x * (long)blockDim.x + threadIdx.x;
    long stride = (long)blockDim.x * gridDim.x;
    long spatial = (long)outD * outH * outW;
    for (long i = idx; i < total; i += stride) {
        int co = (int)((i / spatial) % Cout);
        y[i] = bias[co];
    }
}

__global__ void deconv3d_s2p1d2_k3_kernel(const float* __restrict__ x,
                                          const float* __restrict__ w,
                                          const float* __restrict__ bias,
                                          float* __restrict__ y,
                                          long total_nonzero,
                                          int N, int Cin, int Cout,
                                          int D, int H, int W,
                                          int outD, int outH, int outW,
                                          int has_bias) {
    long idx = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (idx >= total_nonzero) return;

    int iw = (int)(idx % W); idx /= W;
    int ih = (int)(idx % H); idx /= H;
    int id = (int)(idx % D); idx /= D;
    int co = (int)(idx % Cout);
    int n  = (int)(idx / Cout);

    float acc = 0.0f;
    const long xNbase = ((long)n * Cin) * D * H * W;

    // Interior points: all 27 kernel taps are valid.  Output coordinate is (2*id+1,2*ih+1,2*iw+1).
    if (id > 0 && id < D-1 && ih > 0 && ih < H-1 && iw > 0 && iw < W-1) {
        int id0 = id + 1, id1 = id, id2 = id - 1;
        int ih0 = ih + 1, ih1 = ih, ih2 = ih - 1;
        int iw0 = iw + 1, iw1 = iw, iw2 = iw - 1;
        #pragma unroll 1
        for (int ci = 0; ci < Cin; ++ci) {
            const float* xp = x + xNbase + ((long)ci * D * H * W);
            const float* wp = w + (((long)ci * Cout + co) * 27);
            long d0 = (long)id0 * H * W, d1 = (long)id1 * H * W, d2 = (long)id2 * H * W;
            long h0 = (long)ih0 * W,     h1 = (long)ih1 * W,     h2 = (long)ih2 * W;
            acc = fmaf(xp[d0 + h0 + iw0], wp[0], acc);
            acc = fmaf(xp[d0 + h0 + iw1], wp[1], acc);
            acc = fmaf(xp[d0 + h0 + iw2], wp[2], acc);
            acc = fmaf(xp[d0 + h1 + iw0], wp[3], acc);
            acc = fmaf(xp[d0 + h1 + iw1], wp[4], acc);
            acc = fmaf(xp[d0 + h1 + iw2], wp[5], acc);
            acc = fmaf(xp[d0 + h2 + iw0], wp[6], acc);
            acc = fmaf(xp[d0 + h2 + iw1], wp[7], acc);
            acc = fmaf(xp[d0 + h2 + iw2], wp[8], acc);
            acc = fmaf(xp[d1 + h0 + iw0], wp[9], acc);
            acc = fmaf(xp[d1 + h0 + iw1], wp[10], acc);
            acc = fmaf(xp[d1 + h0 + iw2], wp[11], acc);
            acc = fmaf(xp[d1 + h1 + iw0], wp[12], acc);
            acc = fmaf(xp[d1 + h1 + iw1], wp[13], acc);
            acc = fmaf(xp[d1 + h1 + iw2], wp[14], acc);
            acc = fmaf(xp[d1 + h2 + iw0], wp[15], acc);
            acc = fmaf(xp[d1 + h2 + iw1], wp[16], acc);
            acc = fmaf(xp[d1 + h2 + iw2], wp[17], acc);
            acc = fmaf(xp[d2 + h0 + iw0], wp[18], acc);
            acc = fmaf(xp[d2 + h0 + iw1], wp[19], acc);
            acc = fmaf(xp[d2 + h0 + iw2], wp[20], acc);
            acc = fmaf(xp[d2 + h1 + iw0], wp[21], acc);
            acc = fmaf(xp[d2 + h1 + iw1], wp[22], acc);
            acc = fmaf(xp[d2 + h1 + iw2], wp[23], acc);
            acc = fmaf(xp[d2 + h2 + iw0], wp[24], acc);
            acc = fmaf(xp[d2 + h2 + iw1], wp[25], acc);
            acc = fmaf(xp[d2 + h2 + iw2], wp[26], acc);
        }
    } else {
        int jd = id + 1, jh = ih + 1, jw = iw + 1;
        #pragma unroll 1
        for (int ci = 0; ci < Cin; ++ci) {
            const float* xp = x + xNbase + ((long)ci * D * H * W);
            const float* wp = w + (((long)ci * Cout + co) * 27);
            #pragma unroll
            for (int kd = 0; kd < 3; ++kd) {
                int xid = jd - kd;
                if ((unsigned)xid >= (unsigned)D) continue;
                #pragma unroll
                for (int kh = 0; kh < 3; ++kh) {
                    int xih = jh - kh;
                    if ((unsigned)xih >= (unsigned)H) continue;
                    #pragma unroll
                    for (int kw = 0; kw < 3; ++kw) {
                        int xiw = jw - kw;
                        if ((unsigned)xiw >= (unsigned)W) continue;
                        acc = fmaf(xp[((long)xid * H + xih) * W + xiw], wp[kd*9 + kh*3 + kw], acc);
                    }
                }
            }
        }
    }

    long out_index = (((((long)n * Cout + co) * outD + (2*id + 1)) * outH + (2*ih + 1)) * outW + (2*iw + 1));
    y[out_index] = acc + (has_bias ? bias[co] : 0.0f);
}

__global__ void deconv3d_generic_kernel(const float* __restrict__ x,
                                        const float* __restrict__ w,
                                        const float* __restrict__ bias,
                                        float* __restrict__ y,
                                        long total,
                                        int N, int Cin, int Cout,
                                        int D, int H, int W,
                                        int K, int stride, int padding, int dilation,
                                        int outD, int outH, int outW,
                                        int has_bias) {
    long idx0 = blockIdx.x * (long)blockDim.x + threadIdx.x;
    if (idx0 >= total) return;
    long idx = idx0;
    int ow = (int)(idx % outW); idx /= outW;
    int oh = (int)(idx % outH); idx /= outH;
    int od = (int)(idx % outD); idx /= outD;
    int co = (int)(idx % Cout);
    int n  = (int)(idx / Cout);
    float acc = has_bias ? bias[co] : 0.0f;
    const long xNbase = ((long)n * Cin) * D * H * W;
    for (int ci = 0; ci < Cin; ++ci) {
        const float* xp = x + xNbase + (long)ci * D * H * W;
        const float* wp = w + (((long)ci * Cout + co) * K * K * K);
        for (int kd = 0; kd < K; ++kd) {
            int td = od + padding - kd * dilation;
            if (td % stride) continue;
            int xid = td / stride;
            if ((unsigned)xid >= (unsigned)D) continue;
            for (int kh = 0; kh < K; ++kh) {
                int th = oh + padding - kh * dilation;
                if (th % stride) continue;
                int xih = th / stride;
                if ((unsigned)xih >= (unsigned)H) continue;
                for (int kw = 0; kw < K; ++kw) {
                    int tw = ow + padding - kw * dilation;
                    if (tw % stride) continue;
                    int xiw = tw / stride;
                    if ((unsigned)xiw >= (unsigned)W) continue;
                    acc = fmaf(xp[((long)xid * H + xih) * W + xiw], wp[((kd*K + kh)*K + kw)], acc);
                }
            }
        }
    }
    y[idx0] = acc;
}

torch::Tensor deconv3d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                    int64_t stride64, int64_t padding64, int64_t dilation64) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "tensors must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat, "float32 only");
    x = x.contiguous();
    weight = weight.contiguous();
    int N = (int)x.size(0), Cin = (int)x.size(1), D = (int)x.size(2), H = (int)x.size(3), W = (int)x.size(4);
    int Cout = (int)weight.size(1), K = (int)weight.size(2);
    int stride = (int)stride64, padding = (int)padding64, dilation = (int)dilation64;
    int outD = (D - 1) * stride - 2 * padding + dilation * (K - 1) + 1;
    int outH = (H - 1) * stride - 2 * padding + dilation * (K - 1) + 1;
    int outW = (W - 1) * stride - 2 * padding + dilation * (K - 1) + 1;
    auto y = at::empty({N, Cout, outD, outH, outW}, x.options());
    bool has_bias = (bias.defined() && bias.numel() > 0);
    if (has_bias) bias = bias.contiguous();
    const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;
    float* yptr = y.data_ptr<float>();
    auto stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;

    if (stride == 2 && padding == 1 && dilation == 2 && K == 3) {
        long total = y.numel();
        if (has_bias) {
            int blocks_init = (int)std::min<long>((total + threads - 1) / threads, 65535L);
            init_bias_kernel<<<blocks_init, threads, 0, stream>>>(yptr, bptr, total, Cout, outD, outH, outW);
        } else {
            C10_CUDA_CHECK(cudaMemsetAsync(yptr, 0, total * sizeof(float), stream));
        }
        long nz = (long)N * Cout * D * H * W;
        int blocks = (int)((nz + threads - 1) / threads);
        deconv3d_s2p1d2_k3_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bptr, yptr,
                                                                  nz, N, Cin, Cout, D, H, W, outD, outH, outW, has_bias ? 1 : 0);
    } else {
        long total = y.numel();
        int blocks = (int)((total + threads - 1) / threads);
        deconv3d_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bptr, yptr,
                                                               total, N, Cin, Cout, D, H, W, K, stride, padding, dilation,
                                                               outD, outH, outW, has_bias ? 1 : 0);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='deconv3d_sparse_s2p1d2_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['deconv3d_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(
            in_channels, out_channels,
            kernel_size=(kernel_size, kernel_size, kernel_size),
            stride=stride, padding=padding, dilation=dilation, bias=bias
        )
        self.stride = int(stride)
        self.padding = int(padding)
        self.dilation = int(dilation)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv_transpose3d.bias
        if b is None:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.deconv3d_forward(x, self.conv_transpose3d.weight, b, self.stride, self.padding, self.dilation)