```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <cuda_runtime.h>

#define IC 3
#define OC 16
#define PMAX 16

// ---------------------------------------------------------------------------
// Eval mega-fusion:
//   ConvTranspose3d(k=3,s=2,p=1)  ==  zero-upsample(2x) + corr(k=3,pad=1, flipped w)
//   AvgPool(2) x2               ==  AvgPool(4, stride 4)
//   BatchNorm(eval)             ==  per-channel affine -> folded into weights
// After pooling only xu[0 .. 4T+3] matter and odd xu entries are zero, so each
// 4x4x4 block average collapses to 8 box-sums (2x2x2 corners) of the *input*,
// combined with pattern-folded weights (k in {0,1} share the same window).
// out[n,o,T] = sum_{i,p} We[i,o,p] * BoxSum_p(x[n,i], T) + bias2[o]
// ---------------------------------------------------------------------------
__global__ void fused_eval_kernel(const float* __restrict__ x,
                                  const float* __restrict__ We,
                                  const float* __restrict__ bias2,
                                  float* __restrict__ out,
                                  int S, int P, int N) {
    __shared__ float sWe[IC][OC][8];
    __shared__ float sBias[OC];
    for (int i = threadIdx.x; i < IC * OC * 8; i += blockDim.x) {
        sWe[0][0][i] = We[i];
    }
    if (threadIdx.x < OC) sBias[threadIdx.x] = bias2[threadIdx.x];
    __syncthreads();

    long total = (long)N * P * P * P;
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int tw = (int)(idx % P);
    int th = (int)((idx / P) % P);
    int td = (int)((idx / (P * P)) % P);
    int n  = (int)(idx / ((long)P * P * P));

    const int S2 = S * S;
    const float* xn = x + (long)n * IC * S2 * S;

    float acc[OC];
#pragma unroll
    for (int o = 0; o < OC; ++o) acc[o] = sBias[o];

#pragma unroll
    for (int c = 0; c < IC; ++c) {
        const float* xc = xn + (long)c * S2 * S;
        // 3x3x3 neighbourhood: x[2T + a], a in {0,1,2}
        float pp[3][3][2];
#pragma unroll
        for (int a = 0; a < 3; ++a) {
#pragma unroll
            for (int b = 0; b < 3; ++b) {
                const float* row = xc + (long)(2 * td + a) * S2 + (long)(2 * th + b) * S + 2 * tw;
#pragma unroll
                for (int e = 0; e < 2; ++e) pp[a][b][e] = row[e] + row[e + 1];
            }
        }
        float q[3][2][2];
#pragma unroll
        for (int a = 0; a < 3; ++a)
#pragma unroll
            for (int e2 = 0; e2 < 2; ++e2)
#pragma unroll
                for (int e3 = 0; e3 < 2; ++e3)
                    q[a][e2][e3] = pp[a][e2][e3] + pp[a][e2 + 1][e3];
        float rr[2][2][2];
#pragma unroll
        for (int e1 = 0; e1 < 2; ++e1)
#pragma unroll
            for (int e2 = 0; e2 < 2; ++e2)
#pragma unroll
                for (int e3 = 0; e3 < 2; ++e3)
                    rr[e1][e2][e3] = q[e1][e2][e3] + q[e1 + 1][e2][e3];
#pragma unroll
        for (int o = 0; o < OC; ++o) {
#pragma unroll
            for (int e1 = 0; e1 < 2; ++e1)
#pragma unroll
                for (int e2 = 0; e2 < 2; ++e2)
#pragma unroll
                    for (int e3 = 0; e3 < 2; ++e3)
                        acc[o] += sWe[c][o][(e1 << 2) | (e2 << 1) | e3] * rr[e1][e2][e3];
        }
    }

    long P3 = (long)P * P * P;
    float* on = out + (long)n * OC * P3 + ((long)td * P + th) * P + tw;
#pragma unroll
    for (int o = 0; o < OC; ++o) on[(long)o * P3] = acc[o];
}

torch::Tensor fused_eval(torch::Tensor x, torch::Tensor We, torch::Tensor bias2,
                         int64_t S, int64_t P, int64_t N, int64_t OC_) {
    auto out = torch::empty({N, OC_, P, P, P}, x.options());
    long total = (long)N * P * P * P;
    int threads = 256;
    long blocks = (total + threads - 1) / threads;
    fused_eval_kernel<<<(int)blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), We.data_ptr<float>(), bias2.data_ptr<float>(),
        out.data_ptr<float>(), (int)S, (int)P, (int)N);
    return out;
}

// ---------------------------------------------------------------------------
// Training-path kernels (batch statistics + fused normalize/pool/pool)
// ---------------------------------------------------------------------------
__global__ void reduce_channel_kernel(const float* __restrict__ y,
                                      double* __restrict__ psum,
                                      double* __restrict__ psumsq,
                                      long spatial, int C) {
    int nc = blockIdx.x;
    int n = nc / C;
    int c = nc % C;
    const float* base = y + (long)nc * spatial;
    float s = 0.f, sq = 0.f;
    long i = threadIdx.x;
    long stride = blockDim.x;
    long vecN = spatial >> 2;
    const float4* b4 = reinterpret_cast<const float4*>(base);
    for (long j = threadIdx.x; j < vecN; j += stride) {
        float4 v = b4[j];
        s += (v.x + v.y) + (v.z + v.w);
        sq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
    }
    for (long j = (vecN << 2) + threadIdx.x; j < spatial; j += stride) {
        float v = base[j];
        s += v;
        sq += v * v;
    }
    (void)i;
    __shared__ float shs[256];
    __shared__ float shq[256];
    shs[threadIdx.x] = s;
    shq[threadIdx.x] = sq;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (threadIdx.x < off) {
            shs[threadIdx.x] += shs[threadIdx.x + off];
            shq[threadIdx.x] += shq[threadIdx.x + off];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        psum[nc] = (double)shs[0];
        psumsq[nc] = (double)shq[0];
    }
}

__global__ void finalize_stats_kernel(const double* __restrict__ psum,
                                      const double* __restrict__ psumsq,
                                      float* __restrict__ running_mean,
                                      float* __restrict__ running_var,
                                      float* __restrict__ mean_out,
                                      float* __restrict__ invstd_out,
                                      float* __restrict__ scale_out,
                                      float* __restrict__ shift_out,
                                      const float* __restrict__ bn_w,
                                      const float* __restrict__ bn_b,
                                      long M, int N, int C,
                                      float eps, float momentum) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= C) return;
    double s = 0.0, sq = 0.0;
    for (int n = 0; n < N; ++n) {
        s += psum[(long)n * C + c];
        sq += psumsq[(long)n * C + c];
    }
    double mean = s / (double)M;
    double var = sq / (double)M - mean * mean;
    if (var < 0.0) var = 0.0;
    double invstd = 1.0 / sqrt(var + (double)eps);
    double var_unb = (M > 1) ? var * ((double)M / (double)(M - 1)) : var;
    running_mean[c] = (float)((1.0 - (double)momentum) * (double)running_mean[c] + (double)momentum * mean);
    running_var[c] = (float)((1.0 - (double)momentum) * (double)running_var[c] + (double)momentum * var_unb);
    mean_out[c] = (float)mean;
    invstd_out[c] = (float)invstd;
    float w = bn_w[c];
    float scale = w * (float)invstd;
    scale_out[c] = scale;
    shift_out[c] = bn_b[c] - (float)mean * scale;
}

// one thread per output voxel: sum 4x4x4 block (float4 along w), then affine
__global__ void bn_pool_kernel(const float* __restrict__ y,
                               const float* __restrict__ scale,
                               const float* __restrict__ shift,
                               float* __restrict__ out,
                               int C, int D, int P) {
    long total = (long)gridDim.x * blockDim.x;
    long stride = total;
    long P3 = (long)P * P * P;
    long totalOut = (long)C * P3;
    int D2 = D * D;
    for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x; idx < totalOut; idx += stride) {
        int tw = (int)(idx % P);
        int th = (int)((idx / P) % P);
        int td = (int)((idx / (P * P)) % P);
        int c = (int)(idx / P3);
        const float* base = y + ((long)c * D2 * D) + (long)(4 * td) * D2 + (long)(4 * th) * D + 4 * tw;
        float s = 0.f;
#pragma unroll
        for (int a = 0; a < 4; ++a) {
#pragma unroll
            for (int b = 0; b < 4; ++b) {
                float4 v = *reinterpret_cast<const float4*>(base + a * D2 + b * D);
                s += (v.x + v.y) + (v.z + v.w);
            }
        }
        out[idx] = s * scale[c] + shift[c];
    }
}

void bn_stats(torch::Tensor y, torch::Tensor psum, torch::Tensor psumsq,
              torch::Tensor running_mean, torch::Tensor running_var,
              torch::Tensor mean_out, torch::Tensor invstd_out,
              torch::Tensor scale_out, torch::Tensor shift_out,
              torch::Tensor bn_w, torch::Tensor bn_b,
              int64_t N, int64_t C, int64_t spatial, double eps, double momentum) {
    auto stream = at::cuda::getCurrentCUDAStream();
    long M = (long)N * spatial;
    reduce_channel_kernel<<<(int)(N * C), 256, 0, stream>>>(
        y.data_ptr<float>(), psum.data_ptr<double>(), psumsq.data_ptr<double>(),
        spatial, (int)C);
    finalize_stats_kernel<<<1, (int)C, 0, stream>>>(
        psum.data_ptr<double>(), psumsq.data_ptr<double>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        mean_out.data_ptr<float>(), invstd_out.data_ptr<float>(),
        scale_out.data_ptr<float>(), shift_out.data_ptr<float>(),
        bn_w.data_ptr<float>(), bn_b.data_ptr<float>(),
        M, (int)N, (int)C, (float)eps, (float)momentum);
}

torch::Tensor bn_pool(torch::Tensor y, torch::Tensor scale, torch::Tensor shift,
                      int64_t C, int64_t D, int64_t P) {
    long NCT = y.size(0);
    auto out = torch::empty({NCT, C, P, P, P}, y.options());
    long totalOut = (long)NCT * C * P * P * P;
    int threads = 256;
    int blocks = (int)std::min<long>((totalOut + threads - 1) / threads, 65535L * 8L);
    bn_pool_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), scale.data_ptr<float>(), shift.data_ptr<float>(),
        out.data_ptr<float>(), (int)(NCT * C), (int)D, (int)P);
    return out;
}
"""

cpp_source = r"""
#include <torch/extension.h>
torch::Tensor fused_eval(torch::Tensor x, torch::Tensor We, torch::Tensor bias2,
                         int64_t S, int64_t P, int64_t N, int64_t OC_);
void bn_stats(torch::Tensor y, torch::Tensor psum, torch::Tensor psumsq,
              torch::Tensor running_mean, torch::Tensor running_var,
              torch::Tensor mean_out, torch::Tensor invstd_out,
              torch::Tensor scale_out, torch::Tensor shift_out,
              torch::Tensor bn_w, torch::Tensor bn_b,
              int64_t N, int64_t C, int64_t spatial, double eps, double momentum);
torch::Tensor bn_pool(torch::Tensor y, torch::Tensor scale, torch::Tensor shift,
                      int64_t C, int64_t D, int64_t P);
"""

_fused = load_inline(
    name="fused_convT_bn_pool3d",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_eval", "bn_stats", "bn_pool"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    Key ideas:
      * ConvTranspose3d(k=3, s=2, p=1) == 2x zero-upsample + 3x3x3 correlation
        (flipped kernel, pad 1).
      * Two AvgPool3d(2) == one AvgPool3d(4, stride 4) (pooled region [0..59]).
      * Eval mode: BatchNorm is a per-channel affine, folded into the conv
        weights. Because odd upsampled entries are zero, each pooled 4x4x4
        average collapses to eight 2x2x2 box sums of the *input*. The whole
        network becomes a single tiny custom CUDA kernel (fused_eval).
      * Training mode: cuDNN transposed conv + custom fused batch-statistics
        kernel + custom fused normalize/pool kernel (BatchNorm folded into the
        pooling affine, both pools merged into one 4x4x4 average).
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size,
                                                 stride=stride, padding=padding)
        self.batch_norm = nn.BatchNorm3d(out_channels)
        self.avg_pool1 = nn.AvgPool3d(kernel_size=2)
        self.avg_pool2 = nn.AvgPool3d(kernel_size=2)
        self._fast = (in_channels == 3 and out_channels == 16 and
                      self.conv_transpose.kernel_size == (3, 3, 3) and
                      self.conv_transpose.stride == (2, 2, 2) and
                      self.conv_transpose.padding == (1, 1, 1) and
                      self.conv_transpose.output_padding == (0, 0, 0) and
                      self.conv_transpose.dilation == (1, 1, 1) and
                      self.conv_transpose.groups == 1 and
                      self.conv_transpose.bias is not None)

    def _stats_training(self):
        bn = self.batch_norm
        return bool(bn.training) or (bn.running_mean is None and bn.running_var is None)

    def _forward_fast_eval(self, x):
        # whole network in one custom kernel
        N, _, S, _, _ = x.shape
        P = (S - 1) // 2  # output spatial size after both pools
        bn = self.batch_norm
        ct = self.conv_transpose
        wf = ct.weight.flip([2, 3, 4])  # (in, out, 3, 3, 3)
        scale = bn.weight / torch.sqrt(bn.running_var + bn.eps)      # (out,)
        shift = bn.bias - bn.running_mean * scale                    # (out,)
        wf = wf * scale.view(1, -1, 1, 1, 1)
        wsel = [wf.select(2, 0) + wf.select(2, 1), wf.select(2, 2)]
        w1 = torch.stack([wsel[0].select(2, 0) + wsel[0].select(2, 1),
                          wsel[0].select(2, 2)], dim=2)              # (in,out,2,3,3)
        w2 = torch.stack([wsel[1].select(2, 0) + wsel[1].select(2, 1),
                          wsel[1].select(2, 2)], dim=2)              # (in,out,2,3,3)
        wa = torch.stack([w1.select(3, 0) + w1.select(3, 1), w1.select(3, 2)], dim=3)
        wb = torch.stack([w2.select(3, 0) + w2.select(3, 1), w2.select(3, 2)], dim=3)
        We = (torch.stack([wa, wb], dim=2) * (1.0 / 64.0)).contiguous()  # (in,out,2,2,2)
        bias2 = (ct.bias * scale + shift).contiguous()
        return _fused.fused_eval(x, We, bias2, S, P, N, wf.shape[1])

    def _forward_train(self, x):
        bn = self.batch_norm
        y = self.conv_transpose(x)
        N, C, D, H, W = y.shape
        spatial = D * H * W
        dev = y.device
        psum = torch.empty(N * C, dtype=torch.float64, device=dev)
        psumsq = torch.empty(N * C, dtype=torch.float64, device=dev)
        mean = torch.empty(C, dtype=torch.float32, device=dev)
        invstd = torch.empty(C, dtype=torch.float32, device=dev)
        scale = torch.empty(C, dtype=torch.float32, device=dev)
        shift = torch.empty(C, dtype=torch.float32, device=dev)
        _fused.bn_stats(y, psum, psumsq, bn.running_mean, bn.running_var,
                        mean, invstd, scale, shift, bn.weight, bn.bias,
                        N, C, spatial, bn.eps, bn.momentum)
        P = D // 4  # 63 -> 15 after the two 2x2x2 average pools
        return _fused.bn_pool(y, scale, shift, C, D, P)

    def forward(self, x):
        if (self._fast and x.is_cuda and x.dtype == torch.float32
                and self.batch_norm.track_running_stats):
            if self._stats_training():
                return self._forward_train(x)
            return self._forward_fast_eval(x)
        # generic fallback (reference behaviour)
        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = self.avg_pool1(x)
        x = self.avg_pool2(x)
        return x
