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

// Fused kernel: computes activation f(x) = x * tanh(softplus(x))
// and then applies BatchNorm: (f(x)-running_mean)/sqrt(running_var+eps)*gamma+beta.
// Assumes input in NCHW layout.
template <typename scalar_t>
__global__ void fused_act_bn_forward_kernel(
    const scalar_t* __restrict__ input,
    scalar_t* __restrict__ output,
    const scalar_t* __restrict__ weight,       // BN gamma
    const scalar_t* __restrict__ bias,         // BN beta
    const scalar_t* __restrict__ running_mean, // BN running mean per channel
    const scalar_t* __restrict__ running_var,  // BN running var per channel
    const float eps,
    const int numel,
    const int C,
    const int H,
    const int W) {

  int index = blockIdx.x * blockDim.x + threadIdx.x;
  const int spatial = H * W;
  while (index < numel) {
    // Compute channel index from NCHW layout.
    int c = (index / spatial) % C;
    scalar_t x_val = input[index];

    if constexpr (std::is_same<scalar_t, float>::value) {
      // Use fast math intrinsics for float.
      scalar_t sp = logf(1.f + __expf(x_val)); // softplus(x)
      scalar_t t = tanhf(sp);
      scalar_t act = x_val * t;
      // Use __ldg to fetch BN parameters from read-only cache.
      scalar_t mean = __ldg(&running_mean[c]);
      scalar_t var  = __ldg(&running_var[c]);
      scalar_t gamma = __ldg(&weight[c]);
      scalar_t beta = __ldg(&bias[c]);
      scalar_t norm = (act - mean) / sqrtf(var + eps);
      output[index] = norm * gamma + beta;
    } else {
      // For double precision, fallback uses standard math functions.
      scalar_t sp = log(1.0 + exp(x_val));
      scalar_t t = tanh(sp);
      scalar_t act = x_val * t;
      scalar_t mean = running_mean[c];
      scalar_t var  = running_var[c];
      scalar_t gamma = weight[c];
      scalar_t beta = bias[c];
      scalar_t norm = (act - mean) / sqrt(var + eps);
      output[index] = norm * gamma + beta;
    }
    index += blockDim.x * gridDim.x;
  }
}

torch::Tensor fused_act_bn_forward(torch::Tensor input,
                                     torch::Tensor weight,
                                     torch::Tensor bias,
                                     torch::Tensor running_mean,
                                     torch::Tensor running_var,
                                     float eps) {
    // Ensure input is contiguous.
    input = input.contiguous();
    auto output = torch::empty_like(input);
    const int numel = input.numel();

    // Input assumed to be in NCHW format.
    const int N = input.size(0);
    const int C = input.size(1);
    const int H = input.size(2);
    const int W = input.size(3);

    const int blockSize = 256;
    const int numBlocks = (numel + blockSize - 1) / blockSize;

    // Guarantee execution on the correct CUDA device.
    const c10::cuda::CUDAGuard device_guard(input.device());

    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "fused_act_bn_forward", ([&] {
        fused_act_bn_forward_kernel<scalar_t><<<numBlocks, blockSize, 0, c10::cuda::getCurrentCUDAStream()>>>(
            input.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            weight.data_ptr<scalar_t>(),
            bias.data_ptr<scalar_t>(),
            running_mean.data_ptr<scalar_t>(),
            running_var.data_ptr<scalar_t>(),
            eps,
            numel,
            C,
            H,
            W);
    }));
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("fused_act_bn_forward", &fused_act_bn_forward,
          "Fused activation (x * tanh(softplus(x))) and BatchNorm forward (CUDA)");
}