import torch
from torch.utils.cpp_extension import load_inline

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

__global__ void atan_sigmoid_mix_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){
    int b = blockIdx.x;
    int lane = blockIdx.y * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.y;
    int row_start = b * D;
    const float* xr = x + row_start;
    float* yr = y + row_start;
    int aligned = ((((long long)xr & 15LL) == 0) && (((long long)yr & 15LL) == 0) && (((long long)scale & 15LL) == 0) && (((long long)bias & 15LL) == 0) && ((D & 3) == 0));
    if(aligned){
        int D4 = (D / 4) * 4;
        #pragma unroll 4
        for(int i = lane * 4; i < D4; i += stride * 4){
            float4 xv = reinterpret_cast<const float4*>(xr)[i / 4];
            float4 sv = reinterpret_cast<const float4*>(scale)[i / 4];
            float4 bv = reinterpret_cast<const float4*>(bias)[i / 4];
            float4 yv;
            float z0 = fmaf(xv.x, sv.x, bv.x);
            float z1 = fmaf(xv.y, sv.y, bv.y);
            float z2 = fmaf(xv.z, sv.z, bv.z);
            float z3 = fmaf(xv.w, sv.w, bv.w);
            float g0 = 1.0f / (1.0f + expf(-(alpha * atanf(z0) + beta)));
            float g1 = 1.0f / (1.0f + expf(-(alpha * atanf(z1) + beta)));
            float g2 = 1.0f / (1.0f + expf(-(alpha * atanf(z2) + beta)));
            float g3 = 1.0f / (1.0f + expf(-(alpha * atanf(z3) + beta)));
            yv.x = xv.x * g0;
            yv.y = xv.y * g1;
            yv.z = xv.z * g2;
            yv.w = xv.w * g3;
            reinterpret_cast<float4*>(yr)[i / 4] = yv;
        }
        #pragma unroll 4
        for(int i = D4 + lane; i < D; i += stride){
            float z = fmaf(xr[i], scale[i], bias[i]);
            float g = 1.0f / (1.0f + expf(-(alpha * atanf(z) + beta)));
            yr[i] = xr[i] * g;
        }
    } else {
        #pragma unroll 4
        for(int i = lane; i < D; i += stride){
            float z = fmaf(xr[i], scale[i], bias[i]);
            float g = 1.0f / (1.0f + expf(-(alpha * atanf(z) + beta)));
            yr[i] = xr[i] * g;
        }
    }
}

torch::Tensor atan_sigmoid_mix_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, 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 a = alpha.item<float>();
    float be = beta.item<float>();
    int block = 256;
    int gy = max(1, min((D + 4095) / 4096, 8));
    dim3 grid(B, gy);
    atan_sigmoid_mix_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
    return y;
}
"""

cpp_source = """
torch::Tensor atan_sigmoid_mix_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
"""

ops = load_inline(
    name="atan_sigmoid_mix_gate",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["atan_sigmoid_mix_gate_cuda"],
    extra_cuda_cflags=["-O3","--use_fast_math"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float):
        super(ModelNew, self).__init__()
        self.ops = ops
        self.register_buffer("scale", scale)
        self.register_buffer("bias", bias)
        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):
        return self.ops.atan_sigmoid_mix_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
