import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__inline__ __device__ float warpReduceSum(float val) {
    for (int offset = 16; offset > 0; offset /= 2)
        val += __shfl_down_sync(0xffffffff, val, offset);
    return val;
}

__inline__ __device__ float blockReduceSum(float val) {
    static __shared__ float shared[32];
    int lane = threadIdx.x % 32;
    int wid = threadIdx.x / 32;

    val = warpReduceSum(val);

    if (lane == 0) shared[wid] = val;
    __syncthreads();

    val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;

    if (wid == 0) val = warpReduceSum(val);

    return val;
}

__global__ void gate_blend_normalize_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    const float* __restrict__ gate,
    float* __restrict__ output,
    int rows,
    int cols,
    float eps
) {
    int bid = blockIdx.x;
    int tid = threadIdx.x;

    if (bid >= rows) return;

    const float* a_row = a + bid * cols;
    const float* b_row = b + bid * cols;
    const float* gate_row = gate + bid * cols;
    float* out_row = output + bid * cols;

    float sum_sq = 0.0f;
    for (int i = tid; i < cols; i += blockDim.x) {
        float g = gate_row[i];
        float blended = g * a_row[i] + (1.0f - g) * b_row[i];
        sum_sq += blended * blended;
    }

    sum_sq = blockReduceSum(sum_sq);

    __shared__ float inv_norm;
    if (tid == 0) {
        inv_norm = rsqrtf(sum_sq + eps);
    }
    __syncthreads();

    float norm_factor = inv_norm;

    for (int i = tid; i < cols; i += blockDim.x) {
        float g = gate_row[i];
        float blended = g * a_row[i] + (1.0f - g) * b_row[i];
        out_row[i] = blended * norm_factor;
    }
}

torch::Tensor gate_blend_normalize_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor gate) {
    auto output = torch::empty_like(a);

    int cols = a.size(a.dim() - 1);
    int rows = a.numel() / cols;

    int block_size = 256;
    while (block_size < cols && block_size < 1024) {
        block_size *= 2;
    }

    gate_blend_normalize_kernel<<<rows, block_size>>>(
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        gate.data_ptr<float>(),
        output.data_ptr<float>(),
        rows,
        cols,
        1e-12f
    );

    return output;
}
"""

cpp_source = """
torch::Tensor gate_blend_normalize_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor gate);
"""

module = load_inline(
    name="gate_blend_normalize",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["gate_blend_normalize_cuda"],
    verbose=True
)


class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.module = module

    def forward(self, a, b, gate):
        return self.module.gate_blend_normalize_cuda(a, b, gate)