import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs 3D Average Pooling.
    """
    def __init__(self, kernel_size: int, stride: int = None, padding: int = 0):
        """
        Initializes the Average Pooling layer.

        Args:
            kernel_size (int): Size of the kernel to apply pooling.
            stride (int, optional): Stride of the pooling operation. Defaults to None, which uses the kernel size.
            padding (int, optional): Padding to apply before pooling. Defaults to 0.
        """
        super(Model, self).__init__()
        self.avg_pool = nn.AvgPool3d(kernel_size=kernel_size, stride=stride, padding=padding)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Average Pooling to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, channels, depth, height, width).

        Returns:
            torch.Tensor: Output tensor with Average Pooling applied, shape depends on kernel_size, stride and padding.
        """
        return self.avg_pool(x)

batch_size = 16
channels = 32
depth = 128
height = 128
width = 256
kernel_size = 3
stride = 2
padding = 1

def get_inputs():
    x = torch.rand(batch_size, channels, depth, height, width)
    return [x]

def get_init_inputs():
    return [kernel_size, stride, padding]

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

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

torch::Tensor avgpool3d_forward_cuda(torch::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding);

torch::Tensor avgpool3d_forward(torch::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding) {
    return avgpool3d_forward_cuda(x, kernel_size, stride, padding);
}
'''

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

static inline long long ceil_div_ll(long long a, long long b) { return (a + b - 1) / b; }

__global__ void avgpool3d_interior_pair_k3s2p1(
    const float* __restrict__ x, float* __restrict__ y,
    int N, int C, int D, int H, int W, int Do, int Ho, int Wo,
    long long total)
{
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int Wint = Wo - 1;
    int P = (Wint + 1) >> 1;
    int p = (int)(idx % P);
    long long t = idx / P;
    int oh_i = (int)(t % (Ho - 1)); t /= (Ho - 1);
    int od_i = (int)(t % (Do - 1)); t /= (Do - 1);
    int nc = (int)t;

    int od = od_i + 1;
    int oh = oh_i + 1;
    int ow0 = 1 + (p << 1);

    int id0 = (od << 1) - 1;
    int ih0 = (oh << 1) - 1;
    int iw0 = (ow0 << 1) - 1;

    long long base = (((long long)nc * D + id0) * H + ih0) * W + iw0;
    int plane = H * W;
    long long out_base = (((long long)nc * Do + od) * Ho + oh) * Wo + ow0;

    float s0 = 0.0f;
    float s1 = 0.0f;

    if (ow0 + 1 < Wo) {
        #pragma unroll
        for (int dz = 0; dz < 3; ++dz) {
            long long bdz = base + (long long)dz * plane;
            #pragma unroll
            for (int dy = 0; dy < 3; ++dy) {
                const float* ptr = x + bdz + (long long)dy * W;
                float a0 = ptr[0];
                float a1 = ptr[1];
                float a2 = ptr[2];
                float a3 = ptr[3];
                float a4 = ptr[4];
                s0 += a0 + a1 + a2;
                s1 += a2 + a3 + a4;
            }
        }
        const float inv27 = 0.037037037037037035f;
        y[out_base] = s0 * inv27;
        y[out_base + 1] = s1 * inv27;
    } else {
        #pragma unroll
        for (int dz = 0; dz < 3; ++dz) {
            long long bdz = base + (long long)dz * plane;
            #pragma unroll
            for (int dy = 0; dy < 3; ++dy) {
                const float* ptr = x + bdz + (long long)dy * W;
                s0 += ptr[0] + ptr[1] + ptr[2];
            }
        }
        y[out_base] = s0 * 0.037037037037037035f;
    }
}

__device__ __forceinline__ float sum_k3s2p1_general_point(
    const float* __restrict__ x, int nc, int od, int oh, int ow,
    int D, int H, int W)
{
    int id_start = (od << 1) - 1;
    int ih_start = (oh << 1) - 1;
    int iw_start = (ow << 1) - 1;
    float s = 0.0f;
    #pragma unroll
    for (int dz = 0; dz < 3; ++dz) {
        int id = id_start + dz;
        if ((unsigned)id < (unsigned)D) {
            #pragma unroll
            for (int dy = 0; dy < 3; ++dy) {
                int ih = ih_start + dy;
                if ((unsigned)ih < (unsigned)H) {
                    long long base = (((long long)nc * D + id) * H + ih) * W;
                    #pragma unroll
                    for (int dx = 0; dx < 3; ++dx) {
                        int iw = iw_start + dx;
                        if ((unsigned)iw < (unsigned)W) s += x[base + iw];
                    }
                }
            }
        }
    }
    return s * 0.037037037037037035f;
}

__global__ void avgpool3d_boundary_k3s2p1(
    const float* __restrict__ x, float* __restrict__ y,
    int N, int C, int D, int H, int W, int Do, int Ho, int Wo,
    long long seg0, long long seg1, long long total)
{
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int nc, od, oh, ow;
    if (idx < seg0) {
        long long t = idx;
        ow = (int)(t % Wo); t /= Wo;
        oh = (int)(t % Ho); t /= Ho;
        od = 0;
        nc = (int)t;
    } else if (idx < seg0 + seg1) {
        long long t = idx - seg0;
        ow = (int)(t % Wo); t /= Wo;
        od = 1 + (int)(t % (Do - 1)); t /= (Do - 1);
        oh = 0;
        nc = (int)t;
    } else {
        long long t = idx - seg0 - seg1;
        oh = 1 + (int)(t % (Ho - 1)); t /= (Ho - 1);
        od = 1 + (int)(t % (Do - 1)); t /= (Do - 1);
        ow = 0;
        nc = (int)t;
    }
    y[(((long long)nc * Do + od) * Ho + oh) * Wo + ow] = sum_k3s2p1_general_point(x, nc, od, oh, ow, D, H, W);
}

__global__ void avgpool3d_generic_kernel(
    const float* __restrict__ x, float* __restrict__ y,
    int N, int C, int D, int H, int W, int Do, int Ho, int Wo,
    int K, int S, int Padd, float inv_div, long long total)
{
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = (int)(idx % Wo);
    long long t = idx / Wo;
    int oh = (int)(t % Ho); t /= Ho;
    int od = (int)(t % Do); t /= Do;
    int nc = (int)t;
    int id_start = od * S - Padd;
    int ih_start = oh * S - Padd;
    int iw_start = ow * S - Padd;
    float s = 0.0f;
    for (int dz = 0; dz < K; ++dz) {
        int id = id_start + dz;
        if ((unsigned)id < (unsigned)D) {
            for (int dy = 0; dy < K; ++dy) {
                int ih = ih_start + dy;
                if ((unsigned)ih < (unsigned)H) {
                    long long base = (((long long)nc * D + id) * H + ih) * W;
                    for (int dx = 0; dx < K; ++dx) {
                        int iw = iw_start + dx;
                        if ((unsigned)iw < (unsigned)W) s += x[base + iw];
                    }
                }
            }
        }
    }
    y[idx] = s * inv_div;
}

torch::Tensor avgpool3d_forward_cuda(torch::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    auto xc = x.contiguous();
    int N = (int)xc.size(0);
    int C = (int)xc.size(1);
    int D = (int)xc.size(2);
    int H = (int)xc.size(3);
    int W = (int)xc.size(4);
    int K = (int)kernel_size;
    int S = stride > 0 ? (int)stride : K;
    int P = (int)padding;
    int Do = (D + 2 * P - K) / S + 1;
    int Ho = (H + 2 * P - K) / S + 1;
    int Wo = (W + 2 * P - K) / S + 1;
    auto y = torch::empty({N, C, Do, Ho, Wo}, xc.options());
    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* xp = xc.data_ptr<float>();
    float* yp = y.data_ptr<float>();
    if (K == 3 && S == 2 && P == 1 && Do > 1 && Ho > 1 && Wo > 1) {
        int Wint = Wo - 1;
        int pairs = (Wint + 1) / 2;
        long long interior_total = (long long)N * C * (Do - 1) * (Ho - 1) * pairs;
        avgpool3d_interior_pair_k3s2p1<<<(int)ceil_div_ll(interior_total, threads), threads, 0, stream>>>(xp, yp, N, C, D, H, W, Do, Ho, Wo, interior_total);
        long long seg0 = (long long)N * C * Ho * Wo;
        long long seg1 = (long long)N * C * (Do - 1) * Wo;
        long long seg2 = (long long)N * C * (Do - 1) * (Ho - 1);
        long long btotal = seg0 + seg1 + seg2;
        avgpool3d_boundary_k3s2p1<<<(int)ceil_div_ll(btotal, threads), threads, 0, stream>>>(xp, yp, N, C, D, H, W, Do, Ho, Wo, seg0, seg1, btotal);
    } else {
        long long total = (long long)N * C * Do * Ho * Wo;
        float inv = 1.0f / (float)(K * K * K);
        avgpool3d_generic_kernel<<<(int)ceil_div_ll(total, threads), threads, 0, stream>>>(xp, yp, N, C, D, H, W, Do, Ho, Wo, K, S, P, inv, total);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='avgpool3d_ext_final',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['avgpool3d_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, kernel_size: int, stride: int = None, padding: int = 0):
        super(ModelNew, self).__init__()
        self.kernel_size = int(kernel_size)
        self.stride = int(stride) if stride is not None else int(kernel_size)
        self.padding = int(padding)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.avgpool3d_forward(x, self.kernel_size, self.stride, self.padding)