import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, scales the output, applies batch normalization, 
    and then performs global average pooling. 
    """
    def __init__(self, in_channels, out_channels, kernel_size, scale_factor, eps=1e-5, momentum=0.1):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = scale_factor
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.scale_factor
        x = self.batch_norm(x)
        x = self.global_avg_pool(x)
        return x

batch_size = 16
in_channels = 64
out_channels = 128
depth, height, width = 16, 32, 32
kernel_size = 5
scale_factor = 2.0

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, scale_factor]

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 bn_w, torch::Tensor bn_b,
                            torch::Tensor running_mean, torch::Tensor running_var,
                            double scale, double eps, double momentum, bool training);
"""

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

__global__ void reduce_x_kernel(const float* __restrict__ x, float* __restrict__ xs,
                                int N, int Ci, int D, int H, int W) {
    int idx = blockIdx.x; int n = idx / Ci; int c = idx - n * Ci; int S = D*H*W;
    const float* ptr = x + ((n * Ci + c) * S);
    float sum = 0.f;
    for (int t=threadIdx.x; t<S; t+=blockDim.x) sum += ptr[t];
    __shared__ float smem[256]; smem[threadIdx.x]=sum; __syncthreads();
    for (int s=blockDim.x>>1; s>0; s>>=1){ if(threadIdx.x<s) smem[threadIdx.x]+=smem[threadIdx.x+s]; __syncthreads(); }
    if(threadIdx.x==0) xs[idx]=smem[0];
}

__global__ void reduce_w_kernel(const float* __restrict__ w, float* __restrict__ ws, int Ci, int Co, int K) {
    int idx=blockIdx.x; int ci=idx/Co; int co=idx-ci*Co; int K3=K*K*K;
    const float* ptr=w+((ci*Co+co)*K3); float sum=0.f;
    for(int t=threadIdx.x;t<K3;t+=blockDim.x) sum += ptr[t];
    __shared__ float smem[128]; smem[threadIdx.x]=sum; __syncthreads();
    for(int s=blockDim.x>>1;s>0;s>>=1){ if(threadIdx.x<s) smem[threadIdx.x]+=smem[threadIdx.x+s]; __syncthreads(); }
    if(threadIdx.x==0) ws[idx]=smem[0];
}

__global__ void eval_out_kernel(const float* __restrict__ xs, const float* __restrict__ ws,
                                const float* __restrict__ bias, const float* __restrict__ bn_w,
                                const float* __restrict__ bn_b, const float* __restrict__ rm,
                                const float* __restrict__ rv, float* __restrict__ out,
                                int N, int Ci, int Co, int Vout, float scale, float eps) {
    int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx>=N*Co) return;
    int n=idx/Co; int co=idx-n*Co; float dot=0.f; const float* xrow=xs+n*Ci;
    for(int ci=0;ci<Ci;++ci) dot += xrow[ci]*ws[ci*Co+co];
    float mean_conv=(bias?bias[co]:0.f)+dot/(float)Vout;
    float z=mean_conv*scale;
    out[idx]=(z-rm[co])*rsqrtf(rv[co]+eps)*bn_w[co]+bn_b[co];
}

__global__ void reduce_y_kernel(const float* __restrict__ y, float* __restrict__ sums, float* __restrict__ sqs,
                                int N, int Co, int S) {
    int idx = blockIdx.x;
    const float* ptr = y + idx * S;
    float sum=0.f, sq=0.f;
    for(int t=threadIdx.x; t<S; t+=blockDim.x){ float v=ptr[t]; sum += v; sq = fmaf(v, v, sq); }
    __shared__ float sm1[256]; __shared__ float sm2[256];
    sm1[threadIdx.x]=sum; sm2[threadIdx.x]=sq; __syncthreads();
    for(int s=blockDim.x>>1; s>0; s>>=1){ if(threadIdx.x<s){ sm1[threadIdx.x]+=sm1[threadIdx.x+s]; sm2[threadIdx.x]+=sm2[threadIdx.x+s]; } __syncthreads(); }
    if(threadIdx.x==0){ sums[idx]=sm1[0]; sqs[idx]=sm2[0]; }
}

__global__ void train_out_kernel(const float* __restrict__ sums, const float* __restrict__ sqs,
                                 const float* __restrict__ bn_w, const float* __restrict__ bn_b,
                                 float* __restrict__ rm, float* __restrict__ rv,
                                 float* __restrict__ out, int N, int Co, int S,
                                 float scale, float eps, float momentum) {
    int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx>=N*Co) return;
    int n=idx/Co; int co=idx-n*Co;
    float tsum=0.f, tsq=0.f;
    #pragma unroll
    for(int nn=0; nn<16; ++nn){
        if(nn < N){ int j=nn*Co+co; tsum += sums[j]; tsq += sqs[j]; }
    }
    float M = (float)(N*S);
    float mean = tsum / M;
    float var = tsq / M - mean * mean;
    if (var < 0.f) var = 0.f;
    float mean_s = mean * scale;
    float var_s = var * scale * scale;
    float sample_mean_s = (sums[idx] / (float)S) * scale;
    out[idx] = (sample_mean_s - mean_s) * rsqrtf(var_s + eps) * bn_w[co] + bn_b[co];
    if(n==0){
        float unbiased = var_s * (M / (M - 1.0f));
        rm[co] = rm[co] * (1.f - momentum) + momentum * mean_s;
        rv[co] = rv[co] * (1.f - momentum) + momentum * unbiased;
    }
}

torch::Tensor eval_fast(torch::Tensor x, torch::Tensor w, torch::Tensor bias, torch::Tensor bn_w, torch::Tensor bn_b,
                        torch::Tensor running_mean, torch::Tensor running_var, double scale_d, double eps_d) {
    x=x.contiguous(); w=w.contiguous();
    int N=x.size(0), Ci=x.size(1), D=x.size(2), H=x.size(3), W=x.size(4); int Co=w.size(1), K=w.size(2);
    int Vout=(D+K-1)*(H+K-1)*(W+K-1);
    auto xs=torch::empty({N,Ci}, x.options()); auto ws=torch::empty({Ci,Co}, x.options()); auto out=torch::empty({N,Co,1,1,1}, x.options());
    cudaStream_t stream=at::cuda::getCurrentCUDAStream();
    reduce_x_kernel<<<N*Ci,256,0,stream>>>(x.data_ptr<float>(), xs.data_ptr<float>(), N,Ci,D,H,W);
    reduce_w_kernel<<<Ci*Co,128,0,stream>>>(w.data_ptr<float>(), ws.data_ptr<float>(), Ci,Co,K);
    eval_out_kernel<<<(N*Co+255)/256,256,0,stream>>>(xs.data_ptr<float>(), ws.data_ptr<float>(), bias.defined()?bias.data_ptr<float>():nullptr, bn_w.data_ptr<float>(), bn_b.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), out.data_ptr<float>(), N,Ci,Co,Vout,(float)scale_d,(float)eps_d);
    return out;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                            torch::Tensor bn_w, torch::Tensor bn_b,
                            torch::Tensor running_mean, torch::Tensor running_var,
                            double scale_d, double eps_d, double momentum_d, bool training) {
    if(!training) return eval_fast(x,w,bias,bn_w,bn_b,running_mean,running_var,scale_d,eps_d);
    x=x.contiguous(); w=w.contiguous();
    c10::optional<at::Tensor> bb = bias.defined() ? c10::optional<at::Tensor>(bias) : c10::nullopt;
    auto y = at::conv_transpose3d(x, w, bb, {1,1,1}, {0,0,0}, {0,0,0}, 1, {1,1,1});
    int N=y.size(0), Co=y.size(1), S=y.size(2)*y.size(3)*y.size(4);
    auto sums=torch::empty({N,Co}, y.options()); auto sqs=torch::empty({N,Co}, y.options()); auto out=torch::empty({N,Co,1,1,1}, y.options());
    cudaStream_t stream=at::cuda::getCurrentCUDAStream();
    reduce_y_kernel<<<N*Co,256,0,stream>>>(y.data_ptr<float>(), sums.data_ptr<float>(), sqs.data_ptr<float>(), N, Co, S);
    train_out_kernel<<<(N*Co+255)/256,256,0,stream>>>(sums.data_ptr<float>(), sqs.data_ptr<float>(), bn_w.data_ptr<float>(), bn_b.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(), out.data_ptr<float>(), N, Co, S, (float)scale_d, (float)eps_d, (float)momentum_d);
    return out;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, scale_factor, eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = scale_factor
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
    def forward(self, x):
        return _ext.fused_forward(x, self.conv_transpose.weight, self.conv_transpose.bias,
                                  self.batch_norm.weight, self.batch_norm.bias,
                                  self.batch_norm.running_mean, self.batch_norm.running_var,
                                  float(self.scale_factor), float(self.batch_norm.eps),
                                  float(self.batch_norm.momentum), bool(self.training and self.batch_norm.training))