import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a series of operations:
    1. Transposed 3D convolution
    2. Mean pooling (across depth)
    3. Addition
    4. Softmax (across channels)
    5. Tanh activation
    6. Scaling
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, scaling_factor):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.bias = nn.Parameter(torch.randn(1, out_channels, 1, 1, 1))  # Broadcastable bias over channels
        self.scaling_factor = scaling_factor

    def forward(self, x):
        x = self.conv_transpose(x)                            # (B, C, D, H, W)
        x = x.mean(dim=2, keepdim=True)                       # Mean pool over depth dim (D)
        x = x + self.bias                                     # Bias add per channel
        x = torch.softmax(x, dim=1)                           # Softmax over channels
        x = torch.tanh(x)                                     # Nonlinearity
        x = x * self.scaling_factor                           # Scaling
        return x

# === Test config ===
batch_size = 16
in_channels  = 16  
out_channels = 64  
depth = 32; height = width = 128  
kernel_size  = 3
stride       = 1  
padding = 1
scaling_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, stride, padding, scaling_factor]

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

cpp_sources = r'''
torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor bias, double scaling);
'''

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

#define IC 16
#define OC 64
#define KD 3
#define KH 3
#define KW 3

__inline__ __device__ float warp_reduce_sum(float v){
    unsigned mask=0xffffffffu;
    v += __shfl_down_sync(mask, v, 16);
    v += __shfl_down_sync(mask, v, 8);
    v += __shfl_down_sync(mask, v, 4);
    v += __shfl_down_sync(mask, v, 2);
    v += __shfl_down_sync(mask, v, 1);
    return v;
}
__inline__ __device__ float warp_reduce_max(float v){
    unsigned mask=0xffffffffu;
    v = fmaxf(v, __shfl_down_sync(mask, v, 16));
    v = fmaxf(v, __shfl_down_sync(mask, v, 8));
    v = fmaxf(v, __shfl_down_sync(mask, v, 4));
    v = fmaxf(v, __shfl_down_sync(mask, v, 2));
    v = fmaxf(v, __shfl_down_sync(mask, v, 1));
    return v;
}

__global__ void reduce_depth_kernel(const float* __restrict__ x, float* __restrict__ total,
                                    int B, int C, int D, int H, int W) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int N = B * C * H * W;
    if (idx >= N) return;
    int w = idx % W;
    int tmp = idx / W;
    int h = tmp % H;
    tmp /= H;
    int c = tmp % C;
    int b = tmp / C;
    const float* p = x + (((b*C + c)*D)*H + h)*W + w;
    float s = 0.0f;
    #pragma unroll
    for (int d = 0; d < 32; ++d) {
        if (d < D) s += p[d * H * W];
    }
    total[idx] = s;
}

__global__ void conv_softmax_kernel(const float* __restrict__ x,
                                    const float* __restrict__ total,
                                    const float* __restrict__ weight,
                                    const float* __restrict__ conv_bias,
                                    const float* __restrict__ bias,
                                    float* __restrict__ out,
                                    int B, int D, int H, int W,
                                    float invD, float scale) {
    __shared__ float buf[4];
    int tid = threadIdx.x;
    int group = tid >> 6;
    int oc = tid & 63;
    int pix = blockIdx.x * 2 + group;
    int P = B * H * W;
    if (pix >= P) return;
    int lane = tid & 31;
    int warp = tid >> 5;

    int rem = pix;
    int ow = rem % W; rem /= W;
    int oh = rem % H; rem /= H;
    int b = rem;

    float acc = 0.0f;
    #pragma unroll
    for (int ic = 0; ic < IC; ++ic) {
        #pragma unroll
        for (int kh = 0; kh < KH; ++kh) {
            int ih = oh + 1 - kh;
            if ((unsigned)ih >= (unsigned)H) continue;
            #pragma unroll
            for (int kw = 0; kw < KW; ++kw) {
                int iw = ow + 1 - kw;
                if ((unsigned)iw >= (unsigned)W) continue;
                int in2_idx = ((b * IC + ic) * H + ih) * W + iw;
                float t = total[in2_idx];
                const float* px = x + (((b * IC + ic) * D) * H + ih) * W + iw;
                float x0 = px[0];
                float xL = px[(D - 1) * H * W];
                int wbase = (((ic * OC + oc) * KD) * KH + kh) * KW + kw;
                float w0 = weight[wbase];
                float w1 = weight[wbase + KH * KW];
                float w2 = weight[wbase + 2 * KH * KW];
                acc = fmaf(t, (w0 + w1 + w2), acc);
                acc = fmaf(-x0, w0, acc);
                acc = fmaf(-xL, w2, acc);
            }
        }
    }
    float v = conv_bias[oc] + bias[oc] + acc * invD;
    float wm = warp_reduce_max(v);
    if (lane == 0) buf[warp] = wm;
    __syncthreads();
    float m = fmaxf(buf[group * 2], buf[group * 2 + 1]);
    float e = expf(v - m);
    float ws = warp_reduce_sum(e);
    if (lane == 0) buf[warp] = ws;
    __syncthreads();
    float s = buf[group * 2] + buf[group * 2 + 1];
    out[((b * OC + oc) * H + oh) * W + ow] = tanhf(e / s) * scale;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor bias, double scaling) {
    x = x.contiguous();
    weight = weight.contiguous();
    conv_bias = conv_bias.contiguous();
    bias = bias.contiguous();
    int B = x.size(0);
    int C = x.size(1);
    int D = x.size(2);
    int H = x.size(3);
    int W = x.size(4);
    auto total = torch::empty({B, C, H, W}, x.options());
    auto out = torch::empty({B, OC, 1, H, W}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int N = B * C * H * W;
    reduce_depth_kernel<<<(N + 255) / 256, 256, 0, stream>>>(x.data_ptr<float>(), total.data_ptr<float>(), B, C, D, H, W);
    int P = B * H * W;
    conv_softmax_kernel<<<(P + 1) / 2, 128, 0, stream>>>(x.data_ptr<float>(), total.data_ptr<float>(), weight.data_ptr<float>(), conv_bias.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), B, D, H, W, 1.0f / (float)D, (float)scaling);
    return out;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, scaling_factor):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.bias = nn.Parameter(torch.randn(1, out_channels, 1, 1, 1))
        self.scaling_factor = scaling_factor
    def forward(self, x):
        return _ext.fused_forward(x, self.conv_transpose.weight, self.conv_transpose.bias, self.bias, float(self.scaling_factor))