import torch
from torch.utils.cpp_extension import load_inline

source = """
#include <torch/extension.h>

__inline__ __device__ float warpReduceMax(float val){
    for(int offset=16; offset>0; offset>>=1){ float v = __shfl_down_sync(0xffffffff, val, offset); val = val > v ? val : v; }
    return val;
}
__inline__ __device__ float warpReduceMin(float val){
    for(int offset=16; offset>0; offset>>=1){ float v = __shfl_down_sync(0xffffffff, val, offset); val = val < v ? val : v; }
    return val;
}

__global__ void range_gate_affine_kernel(const float* x, const float* scale, const float* bias, float* y, int B, int D, float gamma, float beta){
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warpId = tid >> 5;
    extern __shared__ float sh[];
    float lmax = -INFINITY;
    float lmin = INFINITY;
    for(int i = tid * 4; i < D; i += blockDim.x * 4){
        int base = b * D + i;
        for(int k = 0; k < 4 && (i + k) < D; ++k){
            float z = fmaf(x[base + k], scale[i + k], bias[i + k]);
            lmax = z > lmax ? z : lmax;
            lmin = z < lmin ? z : lmin;
        }
    }
    lmax = warpReduceMax(lmax);
    lmin = warpReduceMin(lmin);
    if(lane == 0){
        sh[warpId] = lmax;
        sh[warpId + ((blockDim.x + 31)>>5)] = lmin;
    }
    __syncthreads();
    float rmax = -INFINITY, rmin = INFINITY;
    if(warpId == 0){
        int WN = ((blockDim.x + 31)>>5);
        float vmax = (lane < WN) ? sh[lane] : -INFINITY;
        float vmin = (lane < WN) ? sh[WN + lane] : INFINITY;
        vmax = warpReduceMax(vmax);
        vmin = warpReduceMin(vmin);
        if(lane == 0){ sh[0] = vmax; sh[1] = vmin; }
    }
    __syncthreads();
    rmax = sh[0]; rmin = sh[1];
    float g = 1.0f / (1.0f + expf(-(gamma * (rmax - rmin) + beta)));
    for(int i = tid * 4; i < D; i += blockDim.x * 4){
        int base = b * D + i;
        for(int k = 0; k < 4 && (i + k) < D; ++k){
            float z = fmaf(x[base + k], scale[i + k], bias[i + k]);
            y[base + k] = z * g;
        }
    }
}

torch::Tensor range_gate_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor gamma, torch::Tensor beta){
    auto xc = x.contiguous();
    auto sc = scale.contiguous();
    auto bc = bias.contiguous();
    auto y = torch::empty_like(xc);
    int B = (int)xc.size(0);
    int D = (int)xc.size(1);
    float g = gamma.item<float>();
    float be = beta.item<float>();
    int block = 1024;
    int grid = B;
    size_t shmem = (size_t)((block + 31)>>5) * 2 * sizeof(float);
    range_gate_affine_kernel<<<grid, block, shmem>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, g, be);
    return y;
}
"""

cpp_source = """
torch::Tensor range_gate_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor gamma, torch::Tensor beta);
"""

ops = load_inline(
    name="range_gate_affine",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["range_gate_affine_cuda"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, scale: torch.Tensor, bias: torch.Tensor, gamma: float, beta: float):
        super(ModelNew, self).__init__()
        self.ops = ops
        self.register_buffer("scale", scale)
        self.register_buffer("bias", bias)
        self.register_buffer("gamma", torch.tensor(float(gamma), dtype=torch.float32))
        self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))

    def forward(self, x: torch.Tensor):
        return self.ops.range_gate_affine_cuda(x, self.scale, self.bias, self.gamma, self.beta)
