import torch
import torch.nn as nn

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

    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 square convolution kernel.
        stride (int or tuple, optional): Stride of the convolution. Defaults to 1.
        padding (int or tuple, optional): Padding applied to the input. Defaults to 0.
        output_padding (int or tuple, optional): Additional size added to one side of each dimension in the output shape. 
                                                  Defaults to 0.
        dilation (int or tuple, optional): Spacing between kernel elements. Defaults to 1.
        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: int, stride: int = 1, padding: int = 0, output_padding: int = 0, 
                 dilation: int = 1, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(in_channels, out_channels, (kernel_size, kernel_size, kernel_size), 
                                                stride=stride, padding=padding, output_padding=output_padding, 
                                                dilation=dilation, groups=groups, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the transposed 3D 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 = 48
out_channels = 24
kernel_size = 3
depth = 96
height = 96
width = 96

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]  # Provide in_channels, out_channels, kernel_size for initialization

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor convt3d_cublas_forward(torch::Tensor x, torch::Tensor weight);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdint.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")

static inline int grid_for(long long n) {
    long long g = (n + 255LL) / 256LL;
    return (int)(g < 65535LL ? g : 65535LL);
}

__global__ void pack_w_kernel(const float* __restrict__ w, float* __restrict__ wp, int IC, int OC) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = 27 * OC * IC;
    if (idx < total) {
        int ic = idx % IC;
        int t = idx / IC;
        int oc = t % OC;
        int sh = t / OC;
        // wp[sh, oc, ic] = weight[ic, oc, kd, kh, kw]
        wp[((sh * OC + oc) * IC) + ic] = w[((long long)ic * OC + oc) * 27 + sh];
    }
}

__global__ void fill_ptrs_kernel(const float* x, const float* wp, float* temp,
                                 unsigned long long* Ap, unsigned long long* Bp, unsigned long long* Cp,
                                 int Bsz, int IC, int OC, long long S) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = 27 * Bsz;
    if (idx < total) {
        int n = idx % Bsz;
        int sh = idx / Bsz;
        Ap[idx] = (unsigned long long)(x + (long long)n * IC * S);
        Bp[idx] = (unsigned long long)(wp + (long long)sh * OC * IC);
        Cp[idx] = (unsigned long long)(temp + (long long)idx * OC * S);
    }
}

__global__ void reduce_kernel(const float* __restrict__ temp, float* __restrict__ y,
                              int Bsz, int OC, int D, int H, int W, int OD, int OH, int OW,
                              long long S, long long total) {
    long long idx = blockIdx.x * blockDim.x + threadIdx.x;
    long long st = (long long)blockDim.x * gridDim.x;
    for (; idx < total; idx += st) {
        int ow = idx % OW;
        long long q = idx / OW;
        int oh = q % OH; q /= OH;
        int od = q % OD; q /= OD;
        int oc = q % OC;
        int n = q / OC;

        float sum = 0.0f;
        int kd0 = od - (D - 1); if (kd0 < 0) kd0 = 0;
        int kd1 = od < 2 ? od : 2;
        int kh0 = oh - (H - 1); if (kh0 < 0) kh0 = 0;
        int kh1 = oh < 2 ? oh : 2;
        int kw0 = ow - (W - 1); if (kw0 < 0) kw0 = 0;
        int kw1 = ow < 2 ? ow : 2;

        for (int kd = kd0; kd <= kd1; ++kd) {
            int id = od - kd;
            for (int kh = kh0; kh <= kh1; ++kh) {
                int ih = oh - kh;
                for (int kw = kw0; kw <= kw1; ++kw) {
                    int iw = ow - kw;
                    int sh = kd * 9 + kh * 3 + kw;
                    long long ti = (((long long)sh * Bsz + n) * OC + oc) * S + ((long long)id * H + ih) * W + iw;
                    sum += temp[ti];
                }
            }
        }
        y[idx] = sum;
    }
}

torch::Tensor convt3d_cublas_forward(torch::Tensor x, torch::Tensor weight) {
    CHECK_CUDA(x); CHECK_CUDA(weight);
    CHECK_CONTIGUOUS(x); CHECK_CONTIGUOUS(weight);
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x fp32");
    TORCH_CHECK(weight.scalar_type() == torch::kFloat32, "w fp32");

    int Bsz = (int)x.size(0);
    int IC = (int)x.size(1);
    int D = (int)x.size(2);
    int H = (int)x.size(3);
    int W = (int)x.size(4);
    int OC = (int)weight.size(1);
    TORCH_CHECK(weight.size(2) == 3 && weight.size(3) == 3 && weight.size(4) == 3, "only k3");

    int OD = D + 2, OH = H + 2, OW = W + 2;
    long long S = (long long)D * H * W;
    int batch = 27 * Bsz;

    auto y = torch::empty({Bsz, OC, OD, OH, OW}, x.options());
    auto temp = torch::empty({batch, OC, D, H, W}, x.options());
    auto wp = torch::empty({27, OC, IC}, x.options());
    auto ptr_opts = torch::TensorOptions().dtype(torch::kInt64).device(x.device());
    auto Ap = torch::empty({batch}, ptr_opts);
    auto Bp = torch::empty({batch}, ptr_opts);
    auto Cp = torch::empty({batch}, ptr_opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    pack_w_kernel<<<grid_for(27LL * OC * IC), 256, 0, stream>>>(weight.data_ptr<float>(), wp.data_ptr<float>(), IC, OC);
    fill_ptrs_kernel<<<grid_for(batch), 256, 0, stream>>>(
        x.data_ptr<float>(), wp.data_ptr<float>(), temp.data_ptr<float>(),
        (unsigned long long*)Ap.data_ptr<int64_t>(), (unsigned long long*)Bp.data_ptr<int64_t>(), (unsigned long long*)Cp.data_ptr<int64_t>(),
        Bsz, IC, OC, S);

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cublasSetStream(handle, stream);
    cublasSetMathMode(handle, CUBLAS_TF32_TENSOR_OP_MATH);
    float alpha = 1.0f, beta = 0.0f;
    cublasStatus_t st = cublasSgemmBatched(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                                           (int)S, OC, IC,
                                           &alpha,
                                           (const float**)Ap.data_ptr<int64_t>(), (int)S,
                                           (const float**)Bp.data_ptr<int64_t>(), IC,
                                           &beta,
                                           (float**)Cp.data_ptr<int64_t>(), (int)S,
                                           batch);
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "cublas failed");

    reduce_kernel<<<grid_for(y.numel()), 256, 0, stream>>>(temp.data_ptr<float>(), y.data_ptr<float>(),
                                                           Bsz, OC, D, H, W, OD, OH, OW, S, y.numel());
    return y;
}
'''

_ext = load_inline(
    name='convt3d_final_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['convt3d_cublas_forward'],
    extra_cuda_cflags=['-O3'],
    extra_cflags=['-O3'],
    extra_ldflags=['-lcublas'],
    verbose=False,
)

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.convt3d_cublas_forward(x.contiguous(), self.conv_transpose3d.weight.contiguous())