import torch
from torch.utils.cpp_extension import load_inline

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

__global__ void power_sigmoid_gate_kernel(const float* x, float* y, long long size, float alpha, float beta, float p) {
    long long idx = blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = blockDim.x * gridDim.x;
    for (long long i = idx; i < size; i += stride) {
        float u = x[i];
        float r = fmaxf(u, 0.0f);
        float v = __fmaf_rn(alpha, u, beta);
        float s = 1.0f / (1.0f + expf(-v));
        float pw;
        if (fabsf(p - 2.0f) < 1e-7f) {
            pw = r * r;
        } else if (fabsf(p - 3.0f) < 1e-7f) {
            pw = r * r * r;
        } else if (fabsf(p - 1.0f) < 1e-7f) {
            pw = r;
        } else {
            pw = powf(r, p);
        }
        y[i] = pw * s;
    }
}

torch::Tensor power_sigmoid_gate_cuda(torch::Tensor x, torch::Tensor alpha, torch::Tensor beta, torch::Tensor p) {
    auto x_contig = x.contiguous();
    auto y = torch::empty_like(x_contig);
    float a = alpha.item().toFloat();
    float b = beta.item().toFloat();
    float pe = p.item().toFloat();
    long long total = x_contig.numel();
    int block = 512;
    long long grid = (total + block - 1) / block;
    grid = grid > 65535 ? 65535 : grid;
    power_sigmoid_gate_kernel<<<(int)grid, block>>>(x_contig.data_ptr<float>(), y.data_ptr<float>(), total, a, b, pe);
    return y;
}
"""

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

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

class ModelNew(torch.nn.Module):
    def __init__(self, alpha: torch.Tensor, beta: torch.Tensor, p: torch.Tensor):
        super(ModelNew, self).__init__()
        self.ops = ops
        self.register_buffer("alpha", alpha)
        self.register_buffer("beta", beta)
        self.register_buffer("p", p)

    def forward(self, x):
        return self.ops.power_sigmoid_gate_cuda(x, self.alpha, self.beta, self.p)
