import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a 3D transposed convolution operation with asymmetric input and kernel, and optional stride.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (tuple of ints): Size of the convolution kernel in the form (kernel_size_depth, kernel_size_height, kernel_size_width).
        stride (tuple of ints, optional): Stride of the convolution in the form (stride_depth, stride_height, stride_width). Defaults to (1, 1, 1).
        padding (tuple of ints, optional): Padding applied to the input in the form (padding_depth, padding_height, padding_width). Defaults to (0, 0, 0).
        output_padding (tuple of ints, optional): Additional size added to one side of the output shape. Defaults to (0, 0, 0).
        groups (int, optional): Number of blocked connections from input channels to output channels. 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: tuple, stride: tuple = (1, 1, 1), padding: tuple = (0, 0, 0), output_padding: tuple = (0, 0, 0), groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding, groups=groups, 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 = 8
in_channels = 32
out_channels = 32
kernel_size = (3, 5, 7)
depth = 12
height = 24
width = 48
stride = (2, 2, 2)
padding = (1, 2, 3)
output_padding = (1, 1, 1)
groups = 4

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, output_padding, groups]

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

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

torch::Tensor conv_transpose3d_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                    int sd, int sh, int sw,
                                    int pd, int ph, int pw,
                                    int opd, int oph, int opw,
                                    int groups);

torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                       int sd, int sh, int sw,
                                       int pd, int ph, int pw,
                                       int opd, int oph, int opw,
                                       int groups) {
    return conv_transpose3d_cuda(x, weight, bias, sd, sh, sw, pd, ph, pw, opd, oph, opw, groups);
}
'''

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

static __forceinline__ __device__ int ceil_div_pos_maybe(int a, int b) {
    // only used with b > 0; if a <= 0 the clamped lower bound is 0
    return (a <= 0) ? 0 : ((a + b - 1) / b);
}

template <bool HAS_BIAS>
__global__ void convt3d_s2_g4_k357_kernel(const float* __restrict__ x,
                                           const float* __restrict__ w,
                                           const float* __restrict__ bias,
                                           float* __restrict__ y,
                                           int total) {
    // Specialized for the benchmark case:
    // N x 32 x 12 x 24 x 48 -> N x 32 x 24 x 48 x 96
    const int W_in = 48, H_in = 24, D_in = 12;
    const int C_in = 32, C_out = 32;
    const int OW = 96, OH = 48, OD = 24;
    const int Cout_per_g = 8, Cin_per_g = 8;
    const int KD = 3, KH = 5, KW = 7;
    const int padD = 1, padH = 2, padW = 3;

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int t = idx;
    int ow = t % OW; t /= OW;
    int oh = t % OH; t /= OH;
    int od = t % OD; t /= OD;
    int oc = t % C_out;
    int n  = t / C_out;

    int g = oc >> 3;
    int ocg = oc & 7;
    int ic_base = g << 3;

    int td = od + padD;
    int th = oh + padH;
    int tw = ow + padW;

    int id0 = ceil_div_pos_maybe(td - (KD - 1), 2);
    int ih0 = ceil_div_pos_maybe(th - (KH - 1), 2);
    int iw0 = ceil_div_pos_maybe(tw - (KW - 1), 2);
    int id1 = td >> 1; if (id1 > D_in - 1) id1 = D_in - 1;
    int ih1 = th >> 1; if (ih1 > H_in - 1) ih1 = H_in - 1;
    int iw1 = tw >> 1; if (iw1 > W_in - 1) iw1 = W_in - 1;

    float acc = HAS_BIAS ? __ldg(bias + oc) : 0.0f;

    #pragma unroll
    for (int icg = 0; icg < Cin_per_g; ++icg) {
        int ic = ic_base + icg;
        const float* x_ic = x + (((n * C_in + ic) * D_in) * H_in) * W_in;
        const float* w_ic = w + (((ic * Cout_per_g + ocg) * KD) * KH) * KW;

        for (int id = id0; id <= id1; ++id) {
            int kd = td - (id << 1);
            const float* x_d = x_ic + (id * H_in) * W_in;
            const float* w_d = w_ic + (kd * KH) * KW;
            for (int ih = ih0; ih <= ih1; ++ih) {
                int kh = th - (ih << 1);
                const float* x_h = x_d + ih * W_in;
                const float* w_h = w_d + kh * KW;
                #pragma unroll 4
                for (int iw = iw0; iw <= iw1; ++iw) {
                    int kw = tw - (iw << 1);
                    acc = fmaf(__ldg(x_h + iw), __ldg(w_h + kw), acc);
                }
            }
        }
    }
    y[idx] = acc;
}

template <bool HAS_BIAS>
__global__ void convt3d_generic_kernel(const float* __restrict__ x,
                                        const float* __restrict__ w,
                                        const float* __restrict__ bias,
                                        float* __restrict__ y,
                                        int total,
                                        int N, int C_in, int D_in, int H_in, int W_in,
                                        int C_out, int OD, int OH, int OW,
                                        int KD, int KH, int KW,
                                        int sd, int sh, int sw,
                                        int pd, int ph, int pw,
                                        int groups) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int tmp = idx;
    int ow = tmp % OW; tmp /= OW;
    int oh = tmp % OH; tmp /= OH;
    int od = tmp % OD; tmp /= OD;
    int oc = tmp % C_out;
    int n = tmp / C_out;

    int cout_pg = C_out / groups;
    int cin_pg = C_in / groups;
    int g = oc / cout_pg;
    int ocg = oc - g * cout_pg;
    int ic_base = g * cin_pg;

    int td = od + pd, th = oh + ph, tw = ow + pw;
    int id0 = ceil_div_pos_maybe(td - (KD - 1), sd);
    int ih0 = ceil_div_pos_maybe(th - (KH - 1), sh);
    int iw0 = ceil_div_pos_maybe(tw - (KW - 1), sw);
    int id1 = td / sd; if (id1 > D_in - 1) id1 = D_in - 1;
    int ih1 = th / sh; if (ih1 > H_in - 1) ih1 = H_in - 1;
    int iw1 = tw / sw; if (iw1 > W_in - 1) iw1 = W_in - 1;

    float acc = HAS_BIAS ? __ldg(bias + oc) : 0.0f;
    for (int icg = 0; icg < cin_pg; ++icg) {
        int ic = ic_base + icg;
        const float* x_ic = x + (((n * C_in + ic) * D_in) * H_in) * W_in;
        const float* w_ic = w + (((ic * cout_pg + ocg) * KD) * KH) * KW;
        for (int id = id0; id <= id1; ++id) {
            int kd = td - id * sd;
            const float* x_d = x_ic + (id * H_in) * W_in;
            const float* w_d = w_ic + (kd * KH) * KW;
            for (int ih = ih0; ih <= ih1; ++ih) {
                int kh = th - ih * sh;
                const float* x_h = x_d + ih * W_in;
                const float* w_h = w_d + kh * KW;
                for (int iw = iw0; iw <= iw1; ++iw) {
                    int kw = tw - iw * sw;
                    acc = fmaf(__ldg(x_h + iw), __ldg(w_h + kw), acc);
                }
            }
        }
    }
    y[idx] = acc;
}

torch::Tensor conv_transpose3d_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                    int sd, int sh, int sw,
                                    int pd, int ph, int pw,
                                    int opd, int oph, int opw,
                                    int groups) {
    x = x.contiguous();
    weight = weight.contiguous();
    if (bias.numel() != 0) bias = bias.contiguous();

    const int N = (int)x.size(0);
    const int C_in = (int)x.size(1);
    const int D_in = (int)x.size(2);
    const int H_in = (int)x.size(3);
    const int W_in = (int)x.size(4);
    const int KD = (int)weight.size(2);
    const int KH = (int)weight.size(3);
    const int KW = (int)weight.size(4);
    const int C_out = (int)weight.size(1) * groups;
    const int OD = (D_in - 1) * sd - 2 * pd + KD + opd;
    const int OH = (H_in - 1) * sh - 2 * ph + KH + oph;
    const int OW = (W_in - 1) * sw - 2 * pw + KW + opw;

    auto y = torch::empty({N, C_out, OD, OH, OW}, x.options());
    int total = N * C_out * OD * OH * OW;
    const int threads = 256;
    const dim3 blocks((total + threads - 1) / threads);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    bool has_bias = (bias.numel() != 0);

    if (C_in == 32 && C_out == 32 && D_in == 12 && H_in == 24 && W_in == 48 &&
        KD == 3 && KH == 5 && KW == 7 && sd == 2 && sh == 2 && sw == 2 &&
        pd == 1 && ph == 2 && pw == 3 && opd == 1 && oph == 1 && opw == 1 &&
        groups == 4 && OD == 24 && OH == 48 && OW == 96) {
        if (has_bias) convt3d_s2_g4_k357_kernel<true><<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), total);
        else convt3d_s2_g4_k357_kernel<false><<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), nullptr, y.data_ptr<float>(), total);
    } else {
        if (has_bias) convt3d_generic_kernel<true><<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), y.data_ptr<float>(), total, N, C_in, D_in, H_in, W_in, C_out, OD, OH, OW, KD, KH, KW, sd, sh, sw, pd, ph, pw, groups);
        else convt3d_generic_kernel<false><<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), nullptr, y.data_ptr<float>(), total, N, C_in, D_in, H_in, W_in, C_out, OD, OH, OW, KD, KH, KW, sd, sh, sw, pd, ph, pw, groups);
    }
    return y;
}
'''

_ext = load_inline(
    name='convt3d_custom_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['conv_transpose3d_forward'],
    extra_cuda_cflags=['-O3'],
    extra_cflags=['-O3'],
    verbose=False,
)

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

    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)
        s = self.stride; p = self.padding; op = self.output_padding
        return _ext.conv_transpose3d_forward(x, self.conv_transpose3d.weight, b,
                                             int(s[0]), int(s[1]), int(s[2]),
                                             int(p[0]), int(p[1]), int(p[2]),
                                             int(op[0]), int(op[1]), int(op[2]),
                                             int(self.groups))