import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, applies LeakyReLU, multiplies by a learnable parameter, 
    applies LeakyReLU again, and performs a max pooling operation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, multiplier_shape):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
        self.multiplier = nn.Parameter(torch.randn(multiplier_shape))
        self.leaky_relu = nn.LeakyReLU(negative_slope=0.2)
        self.max_pool = nn.MaxPool3d(kernel_size=2)

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.leaky_relu(x)
        x = x * self.multiplier
        x = self.leaky_relu(x)
        x = self.max_pool(x)
        return x

batch_size = 16
in_channels = 16
out_channels = 32
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
multiplier_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, output_padding, multiplier_shape]

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_deconv_lrelu_mul_lrelu_pool_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor m);

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor m) {
    return fused_deconv_lrelu_mul_lrelu_pool_cuda(x, w, b, m);
}
'''

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

#define CIN 16
#define COUT 32

__device__ __forceinline__ float lrelu(float x) {
    return x >= 0.0f ? x : 0.2f * x;
}

__global__ void fused_kernel(const float* __restrict__ x,
                             const float* __restrict__ w,
                             const float* __restrict__ b,
                             const float* __restrict__ mult,
                             float* __restrict__ out,
                             int N, int D, int H, int Wd,
                             long long total) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int pw = idx % Wd;
    long long t = idx / Wd;
    int ph = t % H;
    t /= H;
    int pd = t % D;
    t /= D;
    int co = t % COUT;
    int n = t / COUT;

    float acc0 = b[co];
    float acc1 = acc0;
    float acc2 = acc0;
    float acc3 = acc0;
    float acc4 = acc0;
    float acc5 = acc0;
    float acc6 = acc0;
    float acc7 = acc0;

    #pragma unroll
    for (int ci = 0; ci < CIN; ++ci) {
        const float* wp = w + ((ci * COUT + co) * 27);
        const float* xp_base = x + (((n * CIN + ci) * D) * H) * Wd;

        // q=(0,0,0): input[pd,ph,pw], contributes to all 8 output parities with k in {1,2}^3
        float xv = xp_base[(pd * H + ph) * Wd + pw];
        acc0 += xv * wp[(1*3 + 1)*3 + 1];
        acc1 += xv * wp[(1*3 + 1)*3 + 2];
        acc2 += xv * wp[(1*3 + 2)*3 + 1];
        acc3 += xv * wp[(1*3 + 2)*3 + 2];
        acc4 += xv * wp[(2*3 + 1)*3 + 1];
        acc5 += xv * wp[(2*3 + 1)*3 + 2];
        acc6 += xv * wp[(2*3 + 2)*3 + 1];
        acc7 += xv * wp[(2*3 + 2)*3 + 2];

        if (pw + 1 < Wd) {
            xv = xp_base[(pd * H + ph) * Wd + (pw + 1)];
            acc1 += xv * wp[(1*3 + 1)*3 + 0];
            acc3 += xv * wp[(1*3 + 2)*3 + 0];
            acc5 += xv * wp[(2*3 + 1)*3 + 0];
            acc7 += xv * wp[(2*3 + 2)*3 + 0];
        }
        if (ph + 1 < H) {
            xv = xp_base[(pd * H + (ph + 1)) * Wd + pw];
            acc2 += xv * wp[(1*3 + 0)*3 + 1];
            acc3 += xv * wp[(1*3 + 0)*3 + 2];
            acc6 += xv * wp[(2*3 + 0)*3 + 1];
            acc7 += xv * wp[(2*3 + 0)*3 + 2];
        }
        if (ph + 1 < H && pw + 1 < Wd) {
            xv = xp_base[(pd * H + (ph + 1)) * Wd + (pw + 1)];
            acc3 += xv * wp[(1*3 + 0)*3 + 0];
            acc7 += xv * wp[(2*3 + 0)*3 + 0];
        }
        if (pd + 1 < D) {
            const float* xp_d1 = xp_base + ((pd + 1) * H) * Wd;
            xv = xp_d1[ph * Wd + pw];
            acc4 += xv * wp[(0*3 + 1)*3 + 1];
            acc5 += xv * wp[(0*3 + 1)*3 + 2];
            acc6 += xv * wp[(0*3 + 2)*3 + 1];
            acc7 += xv * wp[(0*3 + 2)*3 + 2];
            if (pw + 1 < Wd) {
                xv = xp_d1[ph * Wd + (pw + 1)];
                acc5 += xv * wp[(0*3 + 1)*3 + 0];
                acc7 += xv * wp[(0*3 + 2)*3 + 0];
            }
            if (ph + 1 < H) {
                xv = xp_d1[(ph + 1) * Wd + pw];
                acc6 += xv * wp[(0*3 + 0)*3 + 1];
                acc7 += xv * wp[(0*3 + 0)*3 + 2];
            }
            if (ph + 1 < H && pw + 1 < Wd) {
                xv = xp_d1[(ph + 1) * Wd + (pw + 1)];
                acc7 += xv * wp[(0*3 + 0)*3 + 0];
            }
        }
    }

    float mm = mult[co];
    float z0 = lrelu(lrelu(acc0) * mm);
    float z1 = lrelu(lrelu(acc1) * mm);
    float z2 = lrelu(lrelu(acc2) * mm);
    float z3 = lrelu(lrelu(acc3) * mm);
    float z4 = lrelu(lrelu(acc4) * mm);
    float z5 = lrelu(lrelu(acc5) * mm);
    float z6 = lrelu(lrelu(acc6) * mm);
    float z7 = lrelu(lrelu(acc7) * mm);
    float mx = fmaxf(fmaxf(fmaxf(z0, z1), fmaxf(z2, z3)), fmaxf(fmaxf(z4, z5), fmaxf(z6, z7)));
    out[idx] = mx;
}

torch::Tensor fused_deconv_lrelu_mul_lrelu_pool_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor m) {
    const int N = (int)x.size(0);
    const int D = (int)x.size(2);
    const int H = (int)x.size(3);
    const int Wd = (int)x.size(4);
    auto out = torch::empty({N, COUT, D, H, Wd}, x.options());
    long long total = (long long)N * COUT * D * H * Wd;
    const int threads = 256;
    const int blocks = (int)((total + threads - 1) / threads);
    fused_kernel<<<blocks, threads>>>(x.data_ptr<float>(), w.data_ptr<float>(), b.data_ptr<float>(), m.data_ptr<float>(), out.data_ptr<float>(), N, D, H, Wd, total);
    return out;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, multiplier_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
        self.multiplier = nn.Parameter(torch.randn(multiplier_shape))
        self.leaky_relu = nn.LeakyReLU(negative_slope=0.2)
        self.max_pool = nn.MaxPool3d(kernel_size=2)

    def forward(self, x):
        return _ext.fused_forward(x.contiguous(), self.conv_transpose.weight.contiguous(), self.conv_transpose.bias.contiguous(), self.multiplier.contiguous())