import torch
from torch.utils.cpp_extension import load_inline

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

__global__ void spatial_diff_gate_kernel(const float* x, float* y, int N, int C, int H, int W, float alpha, float beta){
    long long tid = blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = blockDim.x * gridDim.x;
    long long total = (long long)N * C * H * W;
    for(long long i = tid; i < total; i += stride){
        long long w = i % W;
        long long hwc = i / W;
        long long h = hwc % H;
        long long nc = hwc / H;
        long long n = nc / C;
        long long c = nc % C;
        long long base = ((long long)n * C + c) * (long long)H * (long long)W + (long long)h * W;
        float xv = x[i];
        float prev = (w > 0) ? x[base + (w - 1)] : 0.0f;
        float d = (w > 0) ? (xv - prev) : 0.0f;
        float g = 1.0f / (1.0f + expf(-(alpha * d + beta)));
        y[i] = xv * g;
    }
}

torch::Tensor spatial_diff_gate_cuda(torch::Tensor x, torch::Tensor alpha, torch::Tensor beta){
    auto xc = x.contiguous();
    auto y = torch::empty_like(xc);
    int N = (int)xc.size(0);
    int C = (int)xc.size(1);
    int H = (int)xc.size(2);
    int W = (int)xc.size(3);
    float a = alpha.item<float>();
    float b = beta.item<float>();
    int block = 1024;
    int grid = (int)((long long)N * C * H * W / block);
    if(grid < 1) grid = 1; if(grid > 65535) grid = 65535;
    spatial_diff_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W, a, b);
    return y;
}
"""

cpp_source = """
torch::Tensor spatial_diff_gate_cuda(torch::Tensor x, torch::Tensor alpha, torch::Tensor beta);
"""

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

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

    def forward(self, x: torch.Tensor):
        return self.ops.spatial_diff_gate_cuda(x, self.alpha, self.beta)
