import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, followed by a sum, layer normalization, average pooling, and GELU activation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, sum_weight, norm_shape, pool_kernel_size):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
        self.sum_weight = nn.Parameter(torch.tensor(sum_weight))
        self.norm = nn.LayerNorm(norm_shape)
        self.avg_pool = nn.AvgPool3d(kernel_size=pool_kernel_size)
        self.gelu = nn.GELU()

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x + self.sum_weight
        x = self.norm(x)
        x = self.avg_pool(x)
        x = self.gelu(x)
        return x

batch_size = 32
in_channels = 32
out_channels = 64
depth, height, width = 16, 32, 32
kernel_size = (3, 3, 3)
stride = (2, 2, 2)
padding = (1, 1, 1)
output_padding = (1, 1, 1)
sum_weight = 1.0
norm_shape = (out_channels,)
pool_kernel_size = (2, 2, 2)

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, output_padding, sum_weight, norm_shape, pool_kernel_size]

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

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

torch::Tensor fused_post_launcher(torch::Tensor conv, torch::Tensor sum_weight, torch::Tensor norm_weight, torch::Tensor norm_bias);

torch::Tensor forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor sum_weight, torch::Tensor norm_weight, torch::Tensor norm_bias) {
    auto y = torch::conv_transpose3d(x, weight, bias, {2,2,2}, {1,1,1}, {1,1,1}, 1, {1,1,1});
    return fused_post_launcher(y, sum_weight, norm_weight, norm_bias);
}
'''

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

__inline__ __device__ float warp_sum(float v) {
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}

__global__ __launch_bounds__(64, 8)
void ln_pool_gelu_kernel(const float* __restrict__ y,
                         const float* __restrict__ sumw,
                         const float* __restrict__ gamma,
                         const float* __restrict__ beta,
                         float* __restrict__ out,
                         int N, int C, int D, int H, int W) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;
    int OW = W >> 1;
    int OH = H >> 1;
    int OD = D >> 1;
    int b = blockIdx.x;
    int oh = b % OH; b /= OH;
    int od = b % OD; b /= OD;
    int c = b % C; int n = b / C;

    __shared__ float vals[256];
    __shared__ float partial[8];
    __shared__ float mean[4];
    __shared__ float invstd[4];

    float sw = sumw[0];
    float local[4];
    #pragma unroll
    for (int r = 0; r < 4; ++r) {
        int dd = r >> 1;
        int hh = r & 1;
        long long idx = (((((long long)n*C + c)*D + (od*2 + dd))*H + (oh*2 + hh))*W + tid);
        float v = y[idx] + sw;
        local[r] = v;
        vals[r*64 + tid] = v;
        float s = warp_sum(v);
        if (lane == 0) partial[r*2 + warp] = s;
    }
    __syncthreads();
    if (tid < 4) mean[tid] = (partial[tid*2] + partial[tid*2 + 1]) * (1.0f/64.0f);
    __syncthreads();

    #pragma unroll
    for (int r = 0; r < 4; ++r) {
        float d = local[r] - mean[r];
        float s = warp_sum(d*d);
        if (lane == 0) partial[r*2 + warp] = s;
    }
    __syncthreads();
    if (tid < 4) invstd[tid] = rsqrtf((partial[tid*2] + partial[tid*2 + 1]) * (1.0f/64.0f) + 1.0e-5f);
    __syncthreads();

    if (tid < 32) {
        int w0 = tid << 1;
        float acc = 0.0f;
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            float m = mean[r];
            float inv = invstd[r];
            float v0 = (vals[r*64 + w0] - m) * inv;
            float v1 = (vals[r*64 + w0 + 1] - m) * inv;
            acc += v0 * gamma[w0] + beta[w0];
            acc += v1 * gamma[w0+1] + beta[w0+1];
        }
        float z = acc * 0.125f;
        float gelu = 0.5f * z * (1.0f + erff(z * 0.70710678118654752440f));
        long long oidx = (((((long long)n*C + c)*OD + od)*OH + oh)*OW + tid);
        out[oidx] = gelu;
    }
}

torch::Tensor fused_post_launcher(torch::Tensor conv, torch::Tensor sum_weight, torch::Tensor norm_weight, torch::Tensor norm_bias) {
    int N = (int)conv.size(0);
    int C = (int)conv.size(1);
    int D = (int)conv.size(2);
    int H = (int)conv.size(3);
    int W = (int)conv.size(4);
    auto out = torch::empty({N, C, D/2, H/2, W/2}, conv.options());
    int blocks = N * C * (D/2) * (H/2);
    ln_pool_gelu_kernel<<<blocks, 64, 0, at::cuda::getCurrentCUDAStream()>>>(
        conv.data_ptr<float>(), sum_weight.data_ptr<float>(), norm_weight.data_ptr<float>(), norm_bias.data_ptr<float>(), out.data_ptr<float>(), N, C, D, H, W);
    return out;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, sum_weight, norm_shape, pool_kernel_size):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
        self.sum_weight = nn.Parameter(torch.tensor(sum_weight))
        self.norm = nn.LayerNorm(norm_shape)
        self.avg_pool = nn.AvgPool3d(kernel_size=pool_kernel_size)
        self.gelu = nn.GELU()
    def forward(self, x):
        return _ext.forward(x, self.conv_transpose.weight, self.conv_transpose.bias, self.sum_weight, self.norm.weight, self.norm.bias)