import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model that performs a sequence of operations:
        - ConvTranspose3d
        - MaxPool3d
        - Softmax
        - Subtract
        - Swish
        - Max
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, pool_kernel_size, pool_stride, pool_padding):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
        self.max_pool = nn.MaxPool3d(kernel_size=pool_kernel_size, stride=pool_stride, padding=pool_padding)
        self.subtract = nn.Parameter(torch.randn(out_channels)) # Assuming subtraction is element-wise across channels

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.max_pool(x)
        x = torch.softmax(x, dim=1) # Apply softmax across channels (dim=1)
        x = x - self.subtract.view(1, -1, 1, 1, 1) # Subtract across channels
        x = torch.sigmoid(x) * x # Swish activation
        x = torch.max(x, dim=1)[0] # Max pooling across channels
        return x

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

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

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

#ifndef C10_CUDA_KERNEL_LAUNCH_CHECK
#define C10_CUDA_KERNEL_LAUNCH_CHECK() cudaGetLastError()
#endif

// Fused specialized kernel for:
// ConvTranspose3d(Cin=3,Cout=16,k=3,s=2,p=1,op=1) -> MaxPool3d(k=2,s=2)
// -> softmax over channels -> subtract -> swish -> max over channels.
__global__ __launch_bounds__(128)
void fused_ct_pool_softmax_swish_max_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    const float* __restrict__ sub,
    float* __restrict__ out,
    int N, int D, int H, int W)
{
    const int64_t total = (int64_t)N * D * H * W;
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int ow = (int)(idx % W);
    int64_t t = idx / W;
    int oh = (int)(t % H);
    t /= H;
    int od = (int)(t % D);
    int n = (int)(t / D);

    float pooled[16];
    #pragma unroll
    for (int co = 0; co < 16; ++co) pooled[co] = -CUDART_INF_F;

    // Iterate over the eight ConvTranspose output positions in the 2x2x2 pool window.
    #pragma unroll
    for (int rd = 0; rd < 2; ++rd) {
      int did[2]; int dkd[2]; int nd = 0;
      if (rd == 0) { did[0] = od; dkd[0] = 1; nd = 1; }
      else {
        did[0] = od; dkd[0] = 2; nd = 1;
        if (od + 1 < D) { did[1] = od + 1; dkd[1] = 0; nd = 2; }
      }
      #pragma unroll
      for (int rh = 0; rh < 2; ++rh) {
        int hih[2]; int hkh[2]; int nh = 0;
        if (rh == 0) { hih[0] = oh; hkh[0] = 1; nh = 1; }
        else {
          hih[0] = oh; hkh[0] = 2; nh = 1;
          if (oh + 1 < H) { hih[1] = oh + 1; hkh[1] = 0; nh = 2; }
        }
        #pragma unroll
        for (int rw = 0; rw < 2; ++rw) {
          int wiw[2]; int wkw[2]; int nw = 0;
          if (rw == 0) { wiw[0] = ow; wkw[0] = 1; nw = 1; }
          else {
            wiw[0] = ow; wkw[0] = 2; nw = 1;
            if (ow + 1 < W) { wiw[1] = ow + 1; wkw[1] = 0; nw = 2; }
          }

          float vals[16];
          #pragma unroll
          for (int co = 0; co < 16; ++co) vals[co] = b[co];

          #pragma unroll
          for (int ci = 0; ci < 3; ++ci) {
            #pragma unroll
            for (int ad = 0; ad < 2; ++ad) {
              if (ad >= nd) continue;
              const int id = did[ad];
              const int kd = dkd[ad];
              #pragma unroll
              for (int ah = 0; ah < 2; ++ah) {
                if (ah >= nh) continue;
                const int ih = hih[ah];
                const int kh = hkh[ah];
                #pragma unroll
                for (int aw = 0; aw < 2; ++aw) {
                  if (aw >= nw) continue;
                  const int iw = wiw[aw];
                  const int kw = wkw[aw];
                  const float xv = x[(((((int64_t)n * 3 + ci) * D + id) * H + ih) * W + iw)];
                  const int koff = (kd * 3 + kh) * 3 + kw;
                  const float* __restrict__ wp = w + ci * (16 * 27) + koff;
                  #pragma unroll
                  for (int co = 0; co < 16; ++co) {
                    vals[co] = fmaf(xv, wp[co * 27], vals[co]);
                  }
                }
              }
            }
          }
          #pragma unroll
          for (int co = 0; co < 16; ++co) {
            pooled[co] = fmaxf(pooled[co], vals[co]);
          }
        }
      }
    }

    float mx = pooled[0];
    #pragma unroll
    for (int co = 1; co < 16; ++co) mx = fmaxf(mx, pooled[co]);

    float exps[16];
    float sum = 0.0f;
    #pragma unroll
    for (int co = 0; co < 16; ++co) {
      float e = expf(pooled[co] - mx);
      exps[co] = e;
      sum += e;
    }
    float inv_sum = 1.0f / sum;

    float best = -CUDART_INF_F;
    #pragma unroll
    for (int co = 0; co < 16; ++co) {
      float y = exps[co] * inv_sum - sub[co];
      float sw = y / (1.0f + expf(-y));
      best = fmaxf(best, sw);
    }
    out[idx] = best;
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor sub) {
    const int N = (int)x.size(0);
    const int D = (int)x.size(2);
    const int H = (int)x.size(3);
    const int W = (int)x.size(4);
    auto out = torch::empty({N, D, H, W}, x.options());
    const int threads = 128;
    const int64_t total = (int64_t)N * D * H * W;
    const int blocks = (int)((total + threads - 1) / threads);
    fused_ct_pool_softmax_swish_max_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), sub.data_ptr<float>(), out.data_ptr<float>(), N, D, H, W);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_ext = load_inline(
    name='fused_ct3d_pool_softmax_swish_max_ext_v1',
    cpp_sources='',
    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, stride, padding, output_padding, pool_kernel_size, pool_stride, pool_padding):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
        self.max_pool = nn.MaxPool3d(kernel_size=pool_kernel_size, stride=pool_stride, padding=pool_padding)
        self.subtract = nn.Parameter(torch.randn(out_channels))

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