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

// Standard product reduction kernel.
// Assumes the input tensor is reshaped as [outer, red, inner] and the output
// is of shape [outer, inner]. Each thread computes one output element by looping
// over the reduction dimension.
__global__ void prod_reduce_kernel(const float* __restrict__ input,
                                   float* __restrict__ output,
                                   int outer,
                                   int red,
                                   int inner) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    int total = outer * inner;
    for (; index < total; index += gridDim.x * blockDim.x) {
        int i = index / inner;  // Outer index.
        int j = index % inner;  // Inner index.
        float prod = 1.0f;
        int base = i * red * inner + j;
        for (int k = 0; k < red; k++) {
            prod *= input[base + k * inner];
        }
        output[index] = prod;
    }
}

// Unrolled product reduction kernel using pointer arithmetic.
// This kernel is used when the reduction size is divisible by 4. It unrolls the loop
// four iterations at a time, reducing pointer arithmetic overhead.
__global__ void prod_reduce_kernel_unroll(const float* __restrict__ input,
                                          float* __restrict__ output,
                                          int outer,
                                          int red,
                                          int inner) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    int total = outer * inner;
    for (; index < total; index += gridDim.x * blockDim.x) {
        int i = index / inner;
        int j = index % inner;
        float prod = 1.0f;
        int base = i * red * inner + j;
        const float* ptr = input + base;
        // Always unroll by a factor of 4.
        for (int k = 0; k < red; k += 4) {
            prod *= ptr[0] * ptr[inner] * ptr[2 * inner] * ptr[3 * inner];
            ptr += 4 * inner;
        }
        output[index] = prod;
    }
}

torch::Tensor prod_reduce(torch::Tensor input, int reduction_dim) {
    // Ensure input is contiguous.
    input = input.contiguous();
    // Use CUDAGuard to ensure we execute on the proper device.
    c10::cuda::CUDAGuard device_guard(input.device());

    int dims = input.dim();
    TORCH_CHECK(reduction_dim >= 0 && reduction_dim < dims, "Invalid reduction dimension");

    // Reshape the tensor as [outer, red, inner].
    int red = input.size(reduction_dim);
    int outer = 1;
    for (int i = 0; i < reduction_dim; i++)
        outer *= input.size(i);
    int inner = 1;
    for (int i = reduction_dim + 1; i < dims; i++)
        inner *= input.size(i);

    // Allocate output tensor (flattened) of shape [outer * inner].
    auto output = torch::empty({outer * inner}, input.options());

    int total = outer * inner;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    
    // If red is divisible by 4, use the unrolled kernel.
    if (red % 4 == 0) {
        prod_reduce_kernel_unroll<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
            input.data_ptr<float>(),
            output.data_ptr<float>(),
            outer,
            red,
            inner
        );
    } else {
        prod_reduce_kernel<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
            input.data_ptr<float>(),
            output.data_ptr<float>(),
            outer,
            red,
            inner
        );
    }
    
    // Build output shape by removing the reduction dimension.
    std::vector<int64_t> out_shape;
    for (int i = 0; i < dims; i++) {
        if (i == reduction_dim) continue;
        out_shape.push_back(input.size(i));
    }
    return output.view(out_shape);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("prod_reduce", &prod_reduce, "Product reduction over a specified dimension (CUDA)");
}