import torch
from torch.utils.cpp_extension import load_inline

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

__global__ void exp_square_affine_gate_kernel(const float* x, const float* scale, const float* bias, float* y, int B, int D, float alpha, float beta){
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int stride = blockDim.x;
    int row_start = b * D;
    int aligned = ((row_start & 3) == 0);
    if(aligned){
        int D4 = (D / 4) * 4;
        #pragma unroll 4
        for(int i = tid * 4; i < D4; i += stride * 4){
            int base = row_start + i;
            float4 xv = reinterpret_cast<const float4*>(x)[base / 4];
            float4 sv = reinterpret_cast<const float4*>(scale)[i / 4];
            float4 bv = reinterpret_cast<const float4*>(bias)[i / 4];
            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 s0 = expf(z0 * z0);
            float s1 = expf(z1 * z1);
            float s2 = expf(z2 * z2);
            float s3 = expf(z3 * z3);
            float4 yv;
            yv.x = xv.x * (1.0f / (1.0f + expf(-(alpha * s0 + beta))));
            yv.y = xv.y * (1.0f / (1.0f + expf(-(alpha * s1 + beta))));
            yv.z = xv.z * (1.0f / (1.0f + expf(-(alpha * s2 + beta))));
            yv.w = xv.w * (1.0f / (1.0f + expf(-(alpha * s3 + beta))));
            reinterpret_cast<float4*>(y)[base / 4] = yv;
        }
        #pragma unroll 4
        for(int i = D4 + tid; i < D; i += stride){
            int base = row_start + i;
            float z = fmaf(x[base], scale[i], bias[i]);
            float s = expf(z * z);
            float g = 1.0f / (1.0f + expf(-(alpha * s + beta)));
            y[base] = x[base] * g;
        }
    } else {
        #pragma unroll 4
        for(int i = tid; i < D; i += stride){
            int base = row_start + i;
            float z = fmaf(x[base], scale[i], bias[i]);
            float s = expf(z * z);
            float g = 1.0f / (1.0f + expf(-(alpha * s + beta)));
            y[base] = x[base] * g;
        }
    }
}

torch::Tensor exp_square_affine_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 = 1024;
    int grid = B;
    exp_square_affine_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 exp_square_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
"""

ops = load_inline(
    name="exp_square_affine_gate",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["exp_square_affine_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: torch.Tensor):
        return self.ops.exp_square_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
