import torch
from torch.utils.cpp_extension import load_inline

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

__global__ void contrastive_gate_kernel(const float* a, const float* b, float* y, long long total, float alpha, float beta) {
    long long tid = blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = blockDim.x * gridDim.x;
    long long total4 = (total / 4) * 4;
    for (long long i = tid * 4; i < total4; i += stride * 4) {
        float4 av = reinterpret_cast<const float4*>(a)[i / 4];
        float4 bv = reinterpret_cast<const float4*>(b)[i / 4];
        float4 dv;
        dv.x = av.x - bv.x; dv.y = av.y - bv.y; dv.z = av.z - bv.z; dv.w = av.w - bv.w;
        float4 sv;
        sv.x = 1.0f / (1.0f + expf(-(alpha * (av.x + bv.x) + beta)));
        sv.y = 1.0f / (1.0f + expf(-(alpha * (av.y + bv.y) + beta)));
        sv.z = 1.0f / (1.0f + expf(-(alpha * (av.z + bv.z) + beta)));
        sv.w = 1.0f / (1.0f + expf(-(alpha * (av.w + bv.w) + beta)));
        float4 yv;
        yv.x = tanhf(dv.x) * sv.x;
        yv.y = tanhf(dv.y) * sv.y;
        yv.z = tanhf(dv.z) * sv.z;
        yv.w = tanhf(dv.w) * sv.w;
        reinterpret_cast<float4*>(y)[i / 4] = yv;
    }
    for (long long i = total4 + tid; i < total; i += stride) {
        float d = a[i] - b[i];
        float s = 1.0f / (1.0f + expf(-(alpha * (a[i] + b[i]) + beta)));
        y[i] = tanhf(d) * s;
    }
}

torch::Tensor contrastive_gate_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor alpha, torch::Tensor beta) {
    auto ac = a.contiguous();
    auto bc = b.contiguous();
    auto y = torch::empty_like(ac);
    long long total = ac.numel();
    float al = alpha.item<float>();
    float be = beta.item<float>();
    int block = 1024;
    long long grid = (total + block - 1) / block;
    if (grid > 65535) grid = 65535;
    contrastive_gate_kernel<<<(int)grid, block>>>(ac.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), total, al, be);
    return y;
}
"""

cpp_source = """
torch::Tensor contrastive_gate_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor alpha, torch::Tensor beta);
"""

ops = load_inline(
    name="contrastive_gate",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["contrastive_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, a: torch.Tensor, b: torch.Tensor):
        return self.ops.contrastive_gate_cuda(a, b, self.alpha, self.beta)
