import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self, num_input_features: int, num_output_features: int):
        """
        :param num_input_features: The number of input feature maps
        :param num_output_features: The number of output feature maps
        """
        super(Model, self).__init__()
        self.transition = nn.Sequential(
            nn.BatchNorm2d(num_input_features),
            nn.ReLU(inplace=True),
            nn.Conv2d(num_input_features, num_output_features, kernel_size=1, bias=False),
            nn.AvgPool2d(kernel_size=2, stride=2)
        )

    def forward(self, x):
        """
        :param x: Input tensor of shape (batch_size, num_input_features, height, width)
        :return: Downsampled tensor with reduced number of feature maps
        """
        return self.transition(x)

batch_size = 128
num_input_features = 32
num_output_features = 64
height, width = 256, 256

def get_inputs():
    return [torch.rand(batch_size, num_input_features, height, width)]

def get_init_inputs():
    return [num_input_features, num_output_features]

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

_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor trans_forward_cuda(torch::Tensor x,
                                 torch::Tensor bn_weight,
                                 torch::Tensor bn_bias,
                                 torch::Tensor running_mean,
                                 torch::Tensor running_var,
                                 torch::Tensor conv_weight,
                                 bool training,
                                 double momentum,
                                 double eps);

torch::Tensor trans_forward(torch::Tensor x,
                            torch::Tensor bn_weight,
                            torch::Tensor bn_bias,
                            torch::Tensor running_mean,
                            torch::Tensor running_var,
                            torch::Tensor conv_weight,
                            bool training,
                            double momentum,
                            double eps) {
    return trans_forward_cuda(x, bn_weight, bn_bias, running_mean, running_var,
                              conv_weight, training, momentum, eps);
}
'''

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

#define SP_TILE 16
#define CO_TILE 32
#define CIN_FIXED 32
#define BPC 256

__global__ void eval_scale_shift_kernel(const float* __restrict__ gamma,
                                        const float* __restrict__ beta,
                                        const float* __restrict__ rmean,
                                        const float* __restrict__ rvar,
                                        float* __restrict__ scale,
                                        float* __restrict__ shift,
                                        float eps, int C) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c < C) {
        float sc = gamma[c] * rsqrtf(rvar[c] + eps);
        scale[c] = sc;
        shift[c] = beta[c] - rmean[c] * sc;
    }
}

__global__ void reduce_sum_sumsq_kernel(const float* __restrict__ x,
                                        double* __restrict__ psum,
                                        double* __restrict__ psq,
                                        int N, int C, int H, int W) {
    int c = blockIdx.x;
    int b = blockIdx.y;
    long long HW = (long long)H * W;
    long long M = (long long)N * HW;
    double sum = 0.0;
    double sq = 0.0;
    for (long long m = (long long)b * blockDim.x + threadIdx.x; m < M; m += (long long)BPC * blockDim.x) {
        long long n = m / HW;
        long long hw = m - n * HW;
        float v = x[n * (long long)C * HW + c * HW + hw];
        sum += (double)v;
        sq += (double)v * (double)v;
    }
    __shared__ double ssum[256];
    __shared__ double ssq[256];
    ssum[threadIdx.x] = sum;
    ssq[threadIdx.x] = sq;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            ssum[threadIdx.x] += ssum[threadIdx.x + stride];
            ssq[threadIdx.x] += ssq[threadIdx.x + stride];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        psum[c * BPC + b] = ssum[0];
        psq[c * BPC + b] = ssq[0];
    }
}

__global__ void train_scale_shift_kernel(const double* __restrict__ psum,
                                         const double* __restrict__ psq,
                                         const float* __restrict__ gamma,
                                         const float* __restrict__ beta,
                                         float* __restrict__ running_mean,
                                         float* __restrict__ running_var,
                                         float* __restrict__ scale,
                                         float* __restrict__ shift,
                                         int N, int H, int W,
                                         float momentum, float eps, int C) {
    int c = blockIdx.x;
    __shared__ double ssum[256];
    __shared__ double ssq[256];
    double a = (threadIdx.x < BPC) ? psum[c * BPC + threadIdx.x] : 0.0;
    double b = (threadIdx.x < BPC) ? psq[c * BPC + threadIdx.x] : 0.0;
    ssum[threadIdx.x] = a;
    ssq[threadIdx.x] = b;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            ssum[threadIdx.x] += ssum[threadIdx.x + stride];
            ssq[threadIdx.x] += ssq[threadIdx.x + stride];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        double M = (double)N * (double)H * (double)W;
        double mean = ssum[0] / M;
        double var = ssq[0] / M - mean * mean;
        if (var < 0.0) var = 0.0;
        float inv = rsqrtf((float)var + eps);
        float sc = gamma[c] * inv;
        scale[c] = sc;
        shift[c] = beta[c] - (float)mean * sc;
        // PyTorch BatchNorm running_var uses the unbiased variance estimate.
        float unb = (float)(var * (M / (M - 1.0)));
        running_mean[c] = running_mean[c] * (1.0f - momentum) + momentum * (float)mean;
        running_var[c] = running_var[c] * (1.0f - momentum) + momentum * unb;
    }
}

__global__ void fused_pool_conv_kernel(const float* __restrict__ x,
                                       const float* __restrict__ scale,
                                       const float* __restrict__ shift,
                                       const float* __restrict__ w,
                                       float* __restrict__ out,
                                       int N, int H, int W, int Cout) {
    __shared__ float act[SP_TILE * CIN_FIXED];
    __shared__ float wt[CO_TILE * CIN_FIXED];

    const int tx = threadIdx.x; // spatial row within tile: 0..15
    const int ty = threadIdx.y; // output channel within tile: 0..31
    const int tid = ty * SP_TILE + tx;

    const int OH = H >> 1;
    const int OW = W >> 1;
    const int OHW = OH * OW;
    const int S = N * OHW;
    const int spatial_base = blockIdx.x * SP_TILE;
    const int co_base = blockIdx.y * CO_TILE;

    // Load pooled BN+ReLU activations. Layout is act[spatial][ci].
    if (tid < SP_TILE * CIN_FIXED) {
        int ci = tid / SP_TILE;
        int sr = tid - ci * SP_TILE;
        int s = spatial_base + sr;
        float av = 0.0f;
        if (s < S) {
            int n = s / OHW;
            int r = s - n * OHW;
            int oh = r / OW;
            int ow = r - oh * OW;
            long long base = ((long long)n * CIN_FIXED + ci) * (long long)H * W + (long long)(oh << 1) * W + (ow << 1);
            float sc = scale[ci];
            float sh = shift[ci];
            float v0 = fmaf(sc, x[base], sh);          v0 = v0 > 0.0f ? v0 : 0.0f;
            float v1 = fmaf(sc, x[base + 1], sh);      v1 = v1 > 0.0f ? v1 : 0.0f;
            float v2 = fmaf(sc, x[base + W], sh);      v2 = v2 > 0.0f ? v2 : 0.0f;
            float v3 = fmaf(sc, x[base + W + 1], sh);  v3 = v3 > 0.0f ? v3 : 0.0f;
            av = (v0 + v1 + v2 + v3) * 0.25f;
        }
        act[sr * CIN_FIXED + ci] = av;
    }

    // Load a 32x32 weight tile, w is [Cout, 32, 1, 1].
    for (int k = tid; k < CO_TILE * CIN_FIXED; k += SP_TILE * CO_TILE) {
        int lco = k / CIN_FIXED;
        int ci = k - lco * CIN_FIXED;
        int co = co_base + lco;
        wt[lco * CIN_FIXED + ci] = (co < Cout) ? w[co * CIN_FIXED + ci] : 0.0f;
    }
    __syncthreads();

    int s = spatial_base + tx;
    int co = co_base + ty;
    if (s < S && co < Cout) {
        float sum = 0.0f;
        #pragma unroll
        for (int ci = 0; ci < CIN_FIXED; ++ci) {
            sum = fmaf(act[tx * CIN_FIXED + ci], wt[ty * CIN_FIXED + ci], sum);
        }
        int n = s / OHW;
        int r = s - n * OHW;
        out[((long long)n * Cout + co) * (long long)OHW + r] = sum;
    }
}

torch::Tensor trans_forward_cuda(torch::Tensor x,
                                 torch::Tensor bn_weight,
                                 torch::Tensor bn_bias,
                                 torch::Tensor running_mean,
                                 torch::Tensor running_var,
                                 torch::Tensor conv_weight,
                                 bool training,
                                 double momentum,
                                 double eps) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous NCHW");
    TORCH_CHECK(conv_weight.is_contiguous(), "conv weight must be contiguous");
    int N = (int)x.size(0);
    int C = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    int Cout = (int)conv_weight.size(0);
    TORCH_CHECK(C == CIN_FIXED, "this optimized kernel expects 32 input channels");
    TORCH_CHECK((H % 2) == 0 && (W % 2) == 0, "H and W must be even");

    c10::cuda::CUDAGuard device_guard(x.device());
    auto out = torch::empty({N, Cout, H / 2, W / 2}, x.options());
    auto scale = torch::empty({C}, x.options());
    auto shift = torch::empty({C}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (training) {
        auto dopt = x.options().dtype(torch::kFloat64);
        auto psum = torch::empty({C, BPC}, dopt);
        auto psq = torch::empty({C, BPC}, dopt);
        dim3 rg(C, BPC);
        reduce_sum_sumsq_kernel<<<rg, 256, 0, stream>>>(x.data_ptr<float>(), psum.data_ptr<double>(), psq.data_ptr<double>(), N, C, H, W);
        train_scale_shift_kernel<<<C, 256, 0, stream>>>(psum.data_ptr<double>(), psq.data_ptr<double>(),
            bn_weight.data_ptr<float>(), bn_bias.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            scale.data_ptr<float>(), shift.data_ptr<float>(), N, H, W, (float)momentum, (float)eps, C);
    } else {
        eval_scale_shift_kernel<<<1, 64, 0, stream>>>(bn_weight.data_ptr<float>(), bn_bias.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(), scale.data_ptr<float>(), shift.data_ptr<float>(), (float)eps, C);
    }

    int OH = H / 2, OW = W / 2;
    int S = N * OH * OW;
    dim3 block(SP_TILE, CO_TILE);
    dim3 grid((S + SP_TILE - 1) / SP_TILE, (Cout + CO_TILE - 1) / CO_TILE);
    fused_pool_conv_kernel<<<grid, block, 0, stream>>>(x.data_ptr<float>(), scale.data_ptr<float>(), shift.data_ptr<float>(),
        conv_weight.data_ptr<float>(), out.data_ptr<float>(), N, H, W, Cout);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_ext = load_inline(
    name="transition_fused_pool_conv_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["trans_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, num_input_features: int, num_output_features: int):
        super(ModelNew, self).__init__()
        self.transition = nn.Sequential(
            nn.BatchNorm2d(num_input_features),
            nn.ReLU(inplace=True),
            nn.Conv2d(num_input_features, num_output_features, kernel_size=1, bias=False),
            nn.AvgPool2d(kernel_size=2, stride=2),
        )

    def forward(self, x):
        bn = self.transition[0]
        conv = self.transition[2]
        if self.training and bn.track_running_stats:
            bn.num_batches_tracked.add_(1)
        mom = 0.1 if bn.momentum is None else float(bn.momentum)
        return _ext.trans_forward(x, bn.weight, bn.bias, bn.running_mean, bn.running_var,
                                  conv.weight, bool(self.training), mom, float(bn.eps))