#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

#define EPS 1e-5

// Kernel 1: Compute mean and inverse standard deviation over the width dimension
// for each (n, c, d, h) in X, where X is of shape [N, C, D, H, W].
// Use vectorized loads with float4 when possible.
__global__ void compute_norm_stats_kernel(const float* __restrict__ X,
                                          float* __restrict__ mean,
                                          float* __restrict__ inv_std,
                                          int N, int C, int D, int H, int W,
                                          float sum_weight) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C * D * H;
    if (idx < total) {
        // Decode indices: idx corresponds to (n, c, d, h)
        int temp = idx;
        int h = temp % H;
        temp /= H;
        int d = temp % D;
        temp /= D;
        int c = temp % C;
        int n = temp / C;
        int base = (((n * C + c) * D + d) * H + h) * W;
        float sum_val = 0.0f;
        float sum_sq = 0.0f;
        
        if (W % 4 == 0) {
            int vecCount = W / 4;
            const float4* Xvec = reinterpret_cast<const float4*>(X + base);
            for (int i = 0; i < vecCount; i++) {
                float4 v = Xvec[i];
                float s0 = v.x + sum_weight;
                float s1 = v.y + sum_weight;
                float s2 = v.z + sum_weight;
                float s3 = v.w + sum_weight;
                sum_val += s0 + s1 + s2 + s3;
                sum_sq += s0 * s0 + s1 * s1 + s2 * s2 + s3 * s3;
            }
        } else {
            for (int w = 0; w < W; w++) {
                float val = X[base + w] + sum_weight;
                sum_val += val;
                sum_sq += val * val;
            }
        }
        float m = sum_val / W;
        float var = sum_sq / W - m * m;
        mean[idx] = m;
        inv_std[idx] = rsqrtf(var + EPS);
    }
}

// Device function for GELU activation.
__device__ inline float gelu(float x) {
    const float kBeta = 0.7978845608f;
    const float kAlpha = 0.044715f;
    return 0.5f * x * (1.0f + tanhf(kBeta * (x + kAlpha * x * x * x)));
}

// Kernel 2: For each output element after pooling, perform average pooling over a window
// of size (pd, ph, pw) and then apply GELU. Pooling is performed over the spatial dimensions of X.
// Normalization statistics (mean and inv_std) are indexed by (n, c, d, h).
__global__ void pool_gelu_kernel(const float* __restrict__ X,
                                 const float* __restrict__ mean,
                                 const float* __restrict__ inv_std,
                                 float* __restrict__ Out,
                                 int N, int C, int D, int H, int W,
                                 int pd, int ph, int pw,
                                 float sum_weight) {
    int outD = D / pd;
    int outH = H / ph;
    int outW = W / pw;
    
    int total_out = N * C * outD * outH * outW;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total_out) {
        int temp = idx;
        int w_out = temp % outW; temp /= outW;
        int h_out = temp % outH; temp /= outH;
        int d_out = temp % outD; temp /= outD;
        int c = temp % C;
        int n = temp / C;
        
        float sum_pool = 0.0f;
        int count = pd * ph * pw;
        for (int i = 0; i < pd; i++) {
            int d = d_out * pd + i;
            for (int j = 0; j < ph; j++) {
                int h = h_out * ph + j;
                int norm_idx = (((n * C + c) * D + d) * H + h);
                float m = mean[norm_idx];
                float inv = inv_std[norm_idx];
                for (int k = 0; k < pw; k++) {
                    int w = w_out * pw + k;
                    int in_idx = ((((n * C + c) * D + d) * H + h) * W) + w;
                    float norm_val = (X[in_idx] + sum_weight - m) * inv;
                    sum_pool += norm_val;
                }
            }
        }
        float avg = sum_pool / count;
        Out[idx] = gelu(avg);
    }
}

torch::Tensor forward(torch::Tensor X, float sum_weight, c10::IntArrayRef pool_kernel_size) {
    c10::cuda::CUDAGuard device_guard(X.device());
    int N = X.size(0);
    int C = X.size(1);
    int D = X.size(2);
    int H = X.size(3);
    int W = X.size(4);
    
    auto options = X.options();
    auto mean_tensor = torch::empty({N, C, D, H}, options);
    auto inv_std_tensor = torch::empty({N, C, D, H}, options);
    
    int total_norm = N * C * D * H;
    int threads = 256;
    int blocks = (total_norm + threads - 1) / threads;
    compute_norm_stats_kernel<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
        X.data_ptr<float>(),
        mean_tensor.data_ptr<float>(),
        inv_std_tensor.data_ptr<float>(),
        N, C, D, H, W,
        sum_weight
    );
    
    int pd = pool_kernel_size[0];
    int ph = pool_kernel_size[1];
    int pw = pool_kernel_size[2];
    int outD = D / pd;
    int outH = H / ph;
    int outW = W / pw;
    
    auto output = torch::empty({N, C, outD, outH, outW}, options);
    int total_out = N * C * outD * outH * outW;
    blocks = (total_out + threads - 1) / threads;
    pool_gelu_kernel<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
        X.data_ptr<float>(),
        mean_tensor.data_ptr<float>(),
        inv_std_tensor.data_ptr<float>(),
        output.data_ptr<float>(),
        N, C, D, H, W,
        pd, ph, pw,
        sum_weight
    );
    
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Fused addition, layer normalization (over width), average pooling, and GELU activation");
}