import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, scaling, average pooling, bias addition, and scaling.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, scale1, scale2, bias_shape):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.scale1 = nn.Parameter(torch.tensor(scale1))
        self.avg_pool = nn.AvgPool3d(kernel_size=2)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.scale2 = nn.Parameter(torch.tensor(scale2))

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.scale1
        x = self.avg_pool(x)
        x = x + self.bias
        x = x * self.scale2
        return x

batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
scale1 = 0.5
scale2 = 1.0
bias_shape = (out_channels, 1, 1, 1)

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, scale1, scale2, bias_shape]

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

_cpp_src = r'''
#include <torch/extension.h>
void fused_ct3d_pool_launch(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias,
                            torch::Tensor scale1, torch::Tensor bias, torch::Tensor scale2,
                            torch::Tensor eff, torch::Tensor out);

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias,
                            torch::Tensor scale1, torch::Tensor bias, torch::Tensor scale2,
                            torch::Tensor eff) {
    auto out = torch::empty({x.size(0), weight.size(1), 15, 31, 31}, x.options());
    fused_ct3d_pool_launch(x, weight, conv_bias, scale1, bias, scale2, eff, out);
    return out;
}
'''

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

__global__ void precompute_eff_kernel(const float* __restrict__ w, float* __restrict__ eff) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= 384) return;
    int t = idx;
    int m = t & 7; t >>= 3;
    int co = t & 15; int ci = t >> 4;
    int db = (m >> 2) & 1;
    int hb = (m >> 1) & 1;
    int wb = m & 1;
    const float* wp = w + ((ci * 16 + co) * 27);
    float s = 0.0f;
    if (db == 0) {
        if (hb == 0) {
            if (wb == 0) {
                s += wp[13]; s += wp[14]; s += wp[16]; s += wp[17];
                s += wp[22]; s += wp[23]; s += wp[25]; s += wp[26];
            } else { s += wp[12]; s += wp[15]; s += wp[21]; s += wp[24]; }
        } else {
            if (wb == 0) { s += wp[10]; s += wp[11]; s += wp[19]; s += wp[20]; }
            else { s += wp[9]; s += wp[18]; }
        }
    } else {
        if (hb == 0) {
            if (wb == 0) { s += wp[4]; s += wp[5]; s += wp[7]; s += wp[8]; }
            else { s += wp[3]; s += wp[6]; }
        } else {
            if (wb == 0) { s += wp[1]; s += wp[2]; }
            else { s += wp[0]; }
        }
    }
    eff[idx] = s;
}

__global__ void fused_eff_kernel(const float* __restrict__ x,
                                 const float* __restrict__ eff,
                                 const float* __restrict__ cb,
                                 const float* __restrict__ s1p,
                                 const float* __restrict__ b,
                                 const float* __restrict__ s2p,
                                 float* __restrict__ out) {
    const int S = 14415;
    int sp = blockIdx.x * blockDim.x + threadIdx.x;
    if (sp >= S) return;
    int co = blockIdx.y;
    int n = blockIdx.z;

    int pw = sp % 31;
    int q = sp / 31;
    int ph = q % 31;
    int pd = q / 31;

    float acc = 0.0f;
    #pragma unroll
    for (int ci=0; ci<3; ++ci) {
        const float* xb = x + (((n*3 + ci)*16 + pd)*32 + ph)*32 + pw;
        const float* ew = eff + ((ci*16 + co)*8);
        acc = fmaf(xb[0],       ew[0], acc);
        acc = fmaf(xb[1],       ew[1], acc);
        acc = fmaf(xb[32],      ew[2], acc);
        acc = fmaf(xb[33],      ew[3], acc);
        acc = fmaf(xb[1024],    ew[4], acc);
        acc = fmaf(xb[1025],    ew[5], acc);
        acc = fmaf(xb[1056],    ew[6], acc);
        acc = fmaf(xb[1057],    ew[7], acc);
    }
    out[((n*16 + co)*S) + sp] = ((*s1p) * (cb[co] + acc * 0.125f) + b[co]) * (*s2p);
}

void fused_ct3d_pool_launch(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias,
                            torch::Tensor scale1, torch::Tensor bias, torch::Tensor scale2,
                            torch::Tensor eff, torch::Tensor out) {
    precompute_eff_kernel<<<2, 256>>>(weight.data_ptr<float>(), eff.data_ptr<float>());
    dim3 block(256);
    dim3 grid((14415 + 255) / 256, 16, x.size(0));
    fused_eff_kernel<<<grid, block>>>(x.data_ptr<float>(), eff.data_ptr<float>(), conv_bias.data_ptr<float>(),
                                     scale1.data_ptr<float>(), bias.data_ptr<float>(), scale2.data_ptr<float>(),
                                     out.data_ptr<float>());
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, scale1, scale2, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.scale1 = nn.Parameter(torch.tensor(scale1))
        self.avg_pool = nn.AvgPool3d(kernel_size=2)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self.scale2 = nn.Parameter(torch.tensor(scale2))
        self.register_buffer('_eff', torch.empty((in_channels, out_channels, 8), dtype=torch.float32), persistent=False)
    def forward(self, x):
        return _ext.fused_forward(x, self.conv_transpose.weight, self.conv_transpose.bias, self.scale1, self.bias, self.scale2, self._eff)