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

namespace py = pybind11;

// Optimized warp-level reduction using __shfl_down_sync.
__inline__ __device__ float warpReduceSum(float val) {
    // Use full mask for the current warp.
    for (int offset = warpSize / 2; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

// Optimized fused kernel performing global average pooling, bias addition, log-sum-exp and multiplication by 10.
// Input shape: [B, C, H, W], bias: [C], output: [B, 1].
__global__ void fused_forward_kernel(const float* __restrict__ input,
                                       int B, int C, int H, int W,
                                       const float* __restrict__ bias,
                                       float* __restrict__ output) {
    int b = blockIdx.x;                // block per sample
    int tid = threadIdx.x;
    int nThreads = blockDim.x;
    int spatialSize = H * W;
    const int warpSizeVal = 32;
    int warpId = tid / warpSizeVal;
    int lane = tid % warpSizeVal;

    // Precompute base index for the sample.
    int base_index = b * (C * spatialSize);

    // Local register accumulation for each channel.
    // Assuming C is small (e.g., <= 64).
    float local_sum[64];
#pragma unroll
    for (int c = 0; c < C; c++) {
        local_sum[c] = 0.0f;
    }

    // Each thread processes a subset of spatial positions.
    for (int pos = tid; pos < spatialSize; pos += nThreads) {
        // Using __ldg for read-only cache if possible.
        for (int c = 0; c < C; c++) {
#pragma unroll
            local_sum[c] = fmaf(__ldg(&input[base_index + c * spatialSize + pos]), 1.0f, local_sum[c]);
        }
    }

    // Warp-level reduction: reduce within each warp using __shfl_down_sync.
#pragma unroll
    for (int offset = warpSizeVal / 2; offset > 0; offset /= 2) {
#pragma unroll
        for (int c = 0; c < C; c++) {
            local_sum[c] += __shfl_down_sync(0xffffffff, local_sum[c], offset);
        }
    }

    // Shared memory to store one sum per warp.
    int warpCount = (nThreads + warpSizeVal - 1) / warpSizeVal;
    extern __shared__ float sdata[];  // Size: warpCount * C floats.
    if (lane == 0) {
        for (int c = 0; c < C; c++) {
            sdata[warpId * C + c] = local_sum[c];
        }
    }
    __syncthreads();

    // Final reduction: thread 0 reduces the results from each warp.
    if (tid == 0) {
        float final_sum[64];
#pragma unroll
        for (int c = 0; c < C; c++) {
            final_sum[c] = 0.0f;
        }
        for (int w = 0; w < warpCount; w++) {
#pragma unroll
            for (int c = 0; c < C; c++) {
                final_sum[c] += sdata[w * C + c];
            }
        }
        // Compute the global average for each channel and add the bias.
        float v[64];
        float max_val = -INFINITY;
#pragma unroll
        for (int c = 0; c < C; c++) {
            float avg = final_sum[c] / float(spatialSize);
            v[c] = avg + bias[c];
            if (v[c] > max_val) {
                max_val = v[c];
            }
        }
        // Compute log-sum-exp over channels.
        float sum_exp = 0.0f;
#pragma unroll
        for (int c = 0; c < C; c++) {
            sum_exp += expf(v[c] - max_val);
        }
        float lse = max_val + logf(sum_exp);
        // Multiply by 10 and store the result.
        output[b] = lse * 10.0f;
    }
}

// C++ interface function callable from Python.
// Uses c10::cuda::CUDAGuard to ensure we use the proper CUDA device.
torch::Tensor fused_forward(torch::Tensor input, torch::Tensor bias) {
    const c10::cuda::CUDAGuard device_guard(input.device());
    int B = input.size(0);
    int C = input.size(1);
    int H = input.size(2);
    int W = input.size(3);

    // Allocate output tensor with shape [B, 1].
    auto output = torch::empty({B, 1}, input.options());

    // Set block and grid dimensions.
    int threads = 256;
    dim3 blockDim(threads);
    dim3 gridDim(B); // one block per sample.

    int warpCount = (threads + 31) / 32;
    size_t sharedMemSize = warpCount * C * sizeof(float);

    fused_forward_kernel<<<gridDim, blockDim, sharedMemSize, c10::cuda::getCurrentCUDAStream()>>>(
        input.data_ptr<float>(), B, C, H, W,
        bias.data_ptr<float>(), output.data_ptr<float>()
    );

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("fused_forward", &fused_forward, "Fused transposed-conv post-processing (global avg pool, bias addition, log-sum-exp, multiplication) with warp-level reduction and loop unrolling");
}