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

#define EPS 1e-5

// ---------------- MaxPool2d Kernel ----------------
// Performs a 2x2 max pooling on input tensor with shape [N, C, H, W] (H and W even)
// Produces output tensor of shape [N, C, H/2, W/2]
__global__ void max_pool2d_kernel(const float* __restrict__ input,
                                  float* __restrict__ output,
                                  int N, int C, int H, int W) {
  int pooled_H = H / 2;
  int pooled_W = W / 2;
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  int total = N * C * pooled_H * pooled_W;
  if (index < total) {
    // Calculate output indices.
    int pw = index % pooled_W;
    int tmp = index / pooled_W;
    int ph = tmp % pooled_H;
    tmp = tmp / pooled_H;
    int c = tmp % C;
    int n = tmp / C;
    
    int h_start = ph * 2;
    int w_start = pw * 2;
    
    int in_base = ((n * C + c) * H);
    int idx0 = in_base + h_start;
    int idx1 = idx0 + 1;
    int col0 = w_start;
    int col1 = w_start + 1;
    
    float a = input[(idx0 * W) + col0];
    float b = input[(idx0 * W) + col1];
    float c_val = input[(idx1 * W) + col0];
    float d = input[(idx1 * W) + col1];
    
    float max_val = a;
    if (b > max_val) max_val = b;
    if (c_val > max_val) max_val = c_val;
    if (d > max_val) max_val = d;
    
    output[index] = max_val;
  }
}

// ---------------- Fused GroupNorm Kernel ----------------
// This kernel fuses the group normalization computations (mean/variance reduction and normalization)
// It processes one (sample, group) contiguous block of pooled tensor at a time.
// pooled tensor shape: [N, C, pooled_H, pooled_W]
// Each block handles one group, where each group corresponds to contiguous channels in a sample.
__global__ void group_norm_fused_kernel(float* __restrict__ pooled,
                                        int N, int C, int pooled_H, int pooled_W, int num_groups) {
  int total_groups = N * num_groups;
  int group_id = blockIdx.x; // one block per group
  if (group_id >= total_groups) return;
  
  int n = group_id / num_groups;
  int g = group_id % num_groups;
  int channels_per_group = C / num_groups;
  int group_size = channels_per_group * pooled_H * pooled_W;
  
  // Compute the base index for this group in the flattened pooled tensor.
  // For sample n, channels [g * channels_per_group, (g+1)*channels_per_group - 1] are contiguous.
  int base = n * C * pooled_H * pooled_W + g * channels_per_group * pooled_H * pooled_W;
  
  extern __shared__ float sdata[]; // Shared memory for reduction: first half for sum, second half for sumsq.
  int tid = threadIdx.x;
  float partial_sum = 0.0f;
  float partial_sumsq = 0.0f;
  
  // Each thread processes multiple elements by striding over the group.
  for (int i = tid; i < group_size; i += blockDim.x) {
    float val = pooled[base + i];
    partial_sum += val;
    partial_sumsq += val * val;
  }
  
  sdata[tid] = partial_sum;
  sdata[blockDim.x + tid] = partial_sumsq;
  __syncthreads();
  
  // Reduction to compute total sum and sumsq.
  for (unsigned int stride = blockDim.x / 2; stride > 0; stride /= 2) {
    if (tid < stride) {
      sdata[tid] += sdata[tid + stride];
      sdata[blockDim.x + tid] += sdata[blockDim.x + tid + stride];
    }
    __syncthreads();
  }
  
  float mean = sdata[0] / group_size;
  float var = sdata[blockDim.x] / group_size - mean * mean;
  float inv_std = rsqrtf(var + EPS);
  
  // Normalize each element in this group.
  for (int i = tid; i < group_size; i += blockDim.x) {
    int idx = base + i;
    float val = pooled[idx];
    pooled[idx] = (val - mean) * inv_std;
  }
}

// ---------------- Host Function: Fused Forward ----------------
// This function fuses 2x2 max pooling and the fused group normalization.
// Steps:
//  1. Apply max pooling on input tensor x. Input x shape: [N, C, H, W], H and W must be even.
//  2. Apply fused group normalization on the pooled tensor in-place.
torch::Tensor fused_forward(torch::Tensor x, int num_groups) {
  c10::cuda::CUDAGuard device_guard(x.device());
  auto stream = c10::cuda::getCurrentCUDAStream();
  
  // x shape: [N, C, H, W]
  int N = x.size(0);
  int C = x.size(1);
  int H = x.size(2);
  int W = x.size(3);
  TORCH_CHECK(H % 2 == 0, "H must be even for 2x2 max pooling");
  TORCH_CHECK(W % 2 == 0, "W must be even for 2x2 max pooling");
  
  int pooled_H = H / 2;
  int pooled_W = W / 2;
  
  // Allocate tensor for max pooled output.
  auto pooled = torch::empty({N, C, pooled_H, pooled_W}, x.options());
  
  // Launch MaxPool2d kernel.
  int total_pool = N * C * pooled_H * pooled_W;
  int threads = 256;
  int blocks = (total_pool + threads - 1) / threads;
  max_pool2d_kernel<<<blocks, threads, 0, stream>>>(
      x.data_ptr<float>(),
      pooled.data_ptr<float>(),
      N, C, H, W);
  
  // Launch fused GroupNorm kernel.
  int total_groups = N * num_groups;
  int threads_per_block = 256;
  int shared_mem = threads_per_block * 2 * sizeof(float);
  group_norm_fused_kernel<<<total_groups, threads_per_block, shared_mem, stream>>>(
      pooled.data_ptr<float>(), N, C, pooled_H, pooled_W, num_groups);
  
  return pooled;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fused_forward", &fused_forward, "Fused 2x2 max pooling and group normalization (fused kernel)");
}