```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
from torch.utils.cpp_extension import load_inline


def _triple(v):
    if isinstance(v, (tuple, list)):
        assert len(v) == 3
        return (int(v[0]), int(v[1]), int(v[2]))
    return (int(v), int(v), int(v))


_CPP_SRC = r"""
#include <torch/extension.h>
#include <vector>
torch::Tensor convt3d(torch::Tensor x, torch::Tensor w);
std::vector<torch::Tensor> reduce_spatial(torch::Tensor y);
torch::Tensor apply_demean(torch::Tensor y, torch::Tensor sum_nc,
                           torch::Tensor gamma, torch::Tensor var_c,
                           double inv_S, double eps);
"""

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

#define IC {IC}
#define OC {OC}
#define KD {KD}
#define KH {KH}
#define KW {KW}
#define SD {SD}
#define SH {SH}
#define SW {SW}
#define PD {PD}
#define PH {PH}
#define PW {PW}

// ---------------------------------------------------------------------------
// Direct output-stationary transposed 3D convolution (no bias; the bias is a
// per-channel constant that cancels in the subsequent mean subtraction).
// One thread computes all OC output channels of one spatial output voxel.
// Threads are mapped so that whole warps share the (od,oh,ow) parity, which
// makes the tap pattern warp-uniform (no divergence, full lane utilization).
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(128) convt3d_kernel(
        const float* __restrict__ x,
        const float* __restrict__ w,
        float* __restrict__ y,
        int N, int ID, int IH, int IW, int OD, int OH, int OW)
{
    const int AD = (OD + SD - 1) / SD;
    const int AH = (OH + SH - 1) / SH;
    const int AW = (OW + SW - 1) / SW;
    const long total = (long)N * AD * AH * AW * SD * SH * SW;
    long tid = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= total) return;

    int aw = tid % AW;
    int ah = (tid / AW) % AH;
    int ad = (tid / ((long)AW * AH)) % AD;
    long rem = tid / ((long)AW * AH * AD);
    int pw = rem % SW;
    int ph = (rem / SW) % SH;
    int pd = (rem / (SW * SH)) % SD;
    int n  = rem / (SW * SH * SD);

    int od = ad * SD + pd;
    int oh = ah * SH + ph;
    int ow = aw * SW + pw;

    float acc[OC];
    #pragma unroll
    for (int oc = 0; oc < OC; ++oc) acc[oc] = 0.f;

    // ---- gather valid taps (warp-uniform control flow) ----
    int idv[KD], kid[KD], ndv = 0;
    #pragma unroll
    for (int kd = 0; kd < KD; ++kd) {
        int td = od + PD - kd;
        if (td >= 0 && (td % SD) == 0) {
            int id = td / SD;
            if (id < ID) { idv[ndv] = id; kid[ndv] = kd; ++ndv; }
        }
    }
    int ihv[KH], kih[KH], nhv = 0;
    #pragma unroll
    for (int kh = 0; kh < KH; ++kh) {
        int th = oh + PH - kh;
        if (th >= 0 && (th % SH) == 0) {
            int ih = th / SH;
            if (ih < IH) { ihv[nhv] = ih; kih[nhv] = kh; ++nhv; }
        }
    }
    int iwv[KW], kiw[KW], nwv = 0;
    #pragma unroll
    for (int kw = 0; kw < KW; ++kw) {
        int tw = ow + PW - kw;
        if (tw >= 0 && (tw % SW) == 0) {
            int iw = tw / SW;
            if (iw < IW) { iwv[nwv] = iw; kiw[nwv] = kw; ++nwv; }
        }
    }

    const long x_nstride  = (long)IC * ID * IH * IW;
    const long x_icstride = (long)ID * IH * IW;
    const int  KSIZE = KD * KH * KW;

    for (int a = 0; a < ndv; ++a) {
        int id = idv[a], kd = kid[a];
        for (int b = 0; b < nhv; ++b) {
            int ih = ihv[b], kh = kih[b];
            long row = (long)n * x_nstride + ((long)id * IH + ih) * IW;
            int khw = kd * KH + kh;
            for (int c2 = 0; c2 < nwv; ++c2) {
                int iw = iwv[c2], kw = kiw[c2];
                const float* xp = x + row + iw;
                const float* wp = w + (khw * KW + kw);
                #pragma unroll
                for (int ic = 0; ic < IC; ++ic) {
                    float xv = __ldg(xp + (long)ic * x_icstride);
                    const float* wq = wp + (long)ic * (OC * KSIZE);
                    #pragma unroll
                    for (int oc = 0; oc < OC; ++oc) {
                        acc[oc] = fmaf(xv, __ldg(wq + (long)oc * KSIZE), acc[oc]);
                    }
                }
            }
        }
    }

    const long ostride = (long)OD * OH * OW;
    float* yp = y + ((long)n * OC) * ostride + ((long)od * OH + oh) * OW + ow;
    #pragma unroll
    for (int oc = 0; oc < OC; ++oc) {
        yp[(long)oc * ostride] = acc[oc];
    }
}

// ---------------------------------------------------------------------------
// Per-(n,c) spatial sum / sum-of-squares (double precision accumulation).
// Also atomically accumulates per-channel totals over the batch.
// ---------------------------------------------------------------------------
__global__ void reduce_spatial_kernel(const float* __restrict__ y,
                                      double* __restrict__ sum_nc,
                                      double* __restrict__ chan_sum,
                                      double* __restrict__ chan_sumsq,
                                      int S, int C)
{
    int nc = blockIdx.x;
    int c = nc % C;
    const float* base = y + (long)nc * S;
    double s = 0.0, q = 0.0;
    for (int i = threadIdx.x; i < S; i += blockDim.x) {
        double v = (double)base[i];
        s += v;
        q += v * v;
    }
    __shared__ double sh_s[256];
    __shared__ double sh_q[256];
    int t = threadIdx.x;
    sh_s[t] = s;
    sh_q[t] = q;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (t < off) {
            sh_s[t] += sh_s[t + off];
            sh_q[t] += sh_q[t + off];
        }
        __syncthreads();
    }
    if (t == 0) {
        sum_nc[nc] = sh_s[0];
        atomicAdd(&chan_sum[c], sh_s[0]);
        atomicAdd(&chan_sumsq[c], sh_q[0]);
    }
}

// ---------------------------------------------------------------------------
// Fused BatchNorm scale + per-(n,c) spatial mean subtraction.
//   BN(x) - mean_sp(BN(x)) == scale_c * (x - mean_sp(x)),
//   with scale_c = gamma_c / sqrt(var_c + eps).
// The BN shift (beta), the BN mean, and the conv bias all cancel exactly.
// ---------------------------------------------------------------------------
__global__ void apply_demean_kernel(const float* __restrict__ y,
                                    const double* __restrict__ sum_nc,
                                    const float* __restrict__ gamma,
                                    const float* __restrict__ var_c,
                                    float* __restrict__ out,
                                    double inv_S, float eps,
                                    int S, int C)
{
    int nc = blockIdx.y;
    int c = nc % C;
    float scale = gamma[c] / sqrtf(var_c[c] + eps);
    float mean = (float)(sum_nc[nc] * inv_S);
    long base = (long)nc * S;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < S) {
        out[base + i] = scale * (y[base + i] - mean);
    }
}

// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
torch::Tensor convt3d(torch::Tensor x, torch::Tensor w) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "float32 required");
    auto xc = x.contiguous();
    auto wc = w.contiguous();
    int N  = xc.size(0);
    int ID = xc.size(2), IH = xc.size(3), IW = xc.size(4);
    int OD = (ID - 1) * SD - 2 * PD + KD;
    int OH = (IH - 1) * SH - 2 * PH + KH;
    int OW = (IW - 1) * SW - 2 * PW + KW;
    auto y = torch::empty({N, OC, OD, OH, OW}, xc.options());
    const int AD = (OD + SD - 1) / SD;
    const int AH = (OH + SH - 1) / SH;
    const int AW = (OW + SW - 1) / SW;
    long total = (long)N * AD * AH * AW * SD * SH * SW;
    int block = 128;
    long grid = (total + block - 1) / block;
    convt3d_kernel<<<grid, block>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), y.data_ptr<float>(),
        N, ID, IH, IW, OD, OH, OW);
    return y;
}

std::vector<torch::Tensor> reduce_spatial(torch::Tensor y) {
    auto yc = y.contiguous();
    int N = yc.size(0), C = yc.size(1);
    long S = (long)yc.size(2) * yc.size(3) * yc.size(4);
    auto opts = torch::TensorOptions().dtype(torch::kDouble).device(yc.device());
    auto sum_nc     = torch::empty({(long)N * C}, opts);
    auto chan_sum   = torch::zeros({C}, opts);
    auto chan_sumsq = torch::zeros({C}, opts);
    reduce_spatial_kernel<<<N * C, 256>>>(
        yc.data_ptr<float>(),
        sum_nc.data_ptr<double>(),
        chan_sum.data_ptr<double>(),
        chan_sumsq.data_ptr<double>(),
        (int)S, C);
    return {sum_nc, chan_sum, chan_sumsq};
}

torch::Tensor apply_demean(torch::Tensor y, torch::Tensor sum_nc,
                           torch::Tensor gamma, torch::Tensor var_c,
                           double inv_S, double eps) {
    auto yc = y.contiguous();
    int N = yc.size(0), C = yc.size(1);
    long S = (long)yc.size(2) * yc.size(3) * yc.size(4);
    auto out = torch::empty_like(yc);
    dim3 grid((unsigned)((S + 255) / 256), (unsigned)(N * C));
    apply_demean_kernel<<<grid, 256>>>(
        yc.data_ptr<float>(),
        sum_nc.data_ptr<double>(),
        gamma.data_ptr<float>(),
        var_c.data_ptr<float>(),
        out.data_ptr<float>(),
        inv_S, (float)eps, (int)S, C);
    return out;
}
"""


class ModelNew(nn.Module):
    """
    ConvTranspose3d + BatchNorm3d + spatial mean subtraction, implemented with
    custom CUDA kernels.

    Key mathematical simplification (exact, not an approximation):
        BN(x) - mean_spatial(BN(x)) == scale_c * (x - mean_spatial(x))
    where scale_c = gamma_c / sqrt(var_c + eps). The BN affine shift, the BN
    mean and the conv bias are per-channel constants and therefore cancel in
    the mean subtraction. This lets us:
      1. run a custom direct transposed-conv kernel with no bias add,
      2. fuse BN + mean-subtraction into a single memory-bound elementwise
         kernel (one reduction + one apply instead of several full passes).
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride,
                 padding, bias=True):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size,
            stride=stride, padding=padding, bias=bias)
        self.batch_norm = nn.BatchNorm3d(out_channels)

        kd, kh, kw = _triple(kernel_size)
        sd, sh, sw = _triple(stride)
        pd, ph, pw = _triple(padding)

        cuda_src = _CUDA_TEMPLATE.format(
            IC=in_channels, OC=out_channels,
            KD=kd, KH=kh, KW=kw,
            SD=sd, SH=sh, SW=sw,
            PD=pd, PH=ph, PW=pw)
        name = "fused_ct3d_bn_%d_%d_%d%d%d_%d%d%d_%d%d%d" % (
            in_channels, out_channels, kd, kh, kw, sd, sh, sw, pd, ph, pw)
        self._ext = load_inline(
            name=name,
            cpp_sources=_CPP_SRC,
            cuda_sources=cuda_src,
            functions=["convt3d", "reduce_spatial", "apply_demean"],
            verbose=False,
            extra_cuda_cflags=["-O3"],
        )

    def forward(self, x):
        bn = self.batch_norm
        # Fallback to the reference path for unusual configurations.
        if (not x.is_cuda) or x.dtype != torch.float32 or bn.momentum is None \
                or (bn.training and bn.track_running_stats is False
                    and bn.running_var is None) \
                or (not bn.training and bn.track_running_stats
                    and bn.running_var is None):
            out = self.conv_transpose(x)
            out = bn(out)
            return out - torch.mean(out, dim=(2, 3, 4), keepdim=True)

        ext = self._ext
        w = self.conv_transpose.weight
        y = ext.convt3d(x, w)  # conv bias intentionally omitted (cancels)

        N, C = y.size(0), y.size(1)
        S = float(y.size(2) * y.size(3) * y.size(4))
        sum_nc, chan_sum, chan_sumsq = ext.reduce_spatial(y)

        if bn.training:
            m = float(N) * S
            chan_mean = chan_sum / m
            chan_var = chan_sumsq / m - chan_mean * chan_mean
            chan_var = chan_var.clamp_(min=0.0)
            var_c = chan_var.float()
            if bn.track_running_stats:
                with torch.no_grad():
                    mom = float(bn.momentum)
                    bmean = chan_mean
                    cb = self.conv_transpose.bias
                    if cb is not None:
                        bmean = bmean + cb.detach().to(torch.float64)
                    bn.running_mean.lerp_(bmean.float(), mom)
                    var_unbiased = chan_var * (m / (m - 1.0))
                    bn.running_var.lerp_(var_unbiased.float(), mom)
                    bn.num_batches_tracked.add_(1)
        else:
            var_c = bn.running_var

        gamma = bn.weight
        if gamma is None:
            gamma = torch.ones(C, device=y.device, dtype=torch.float32)
        out = ext.apply_demean(
            y, sum_nc, gamma.contiguous(), var_c.detach().contiguous(),
            1.0 / S, float(bn.eps))
        return out
```