import torch
from torch.utils.cpp_extension import load_inline

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

__inline__ __device__ float warpReduceSum(float val){
    for(int offset=16; offset>0; offset>>=1){
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__global__ void proto_cos_gate_kernel(const float* x, const float* p, float* y, int B, int D, float alpha, float beta){
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warpId = tid >> 5;
    extern __shared__ float sh[];
    float s_dot = 0.0f;
    float s_x2 = 0.0f;
    float s_p2 = 0.0f;
    for(int i = tid; i < D; i += blockDim.x){
        float xv = x[b * D + i];
        float pv = p[i];
        s_dot += xv * pv;
        s_x2 += xv * xv;
        s_p2 += pv * pv;
    }
    s_dot = warpReduceSum(s_dot);
    s_x2 = warpReduceSum(s_x2);
    s_p2 = warpReduceSum(s_p2);
    if(lane == 0){
        sh[warpId] = s_dot;
        sh[warpId + ((blockDim.x + 31)>>5)] = s_x2;
        sh[warpId + 2*((blockDim.x + 31)>>5)] = s_p2;
    }
    __syncthreads();
    float dot = 0.0f, x2 = 0.0f, p2 = 0.0f;
    if(warpId == 0){
        int WN = ((blockDim.x + 31)>>5);
        float vd = (lane < WN) ? sh[lane] : 0.0f;
        float vx = (lane < WN) ? sh[WN + lane] : 0.0f;
        float vp = (lane < WN) ? sh[2*WN + lane] : 0.0f;
        vd = warpReduceSum(vd);
        vx = warpReduceSum(vx);
        vp = warpReduceSum(vp);
        if(lane == 0){
            sh[0] = vd;
            sh[1] = vx;
            sh[2] = vp;
        }
    }
    __syncthreads();
    dot = sh[0]; x2 = sh[1]; p2 = sh[2];
    float nx = sqrtf(x2 + 1e-8f);
    float np = sqrtf(p2 + 1e-8f);
    float cosv = dot / (nx * np);
    float g = 1.0f / (1.0f + expf(-(alpha * cosv + beta)));
    for(int i = tid; i < D; i += blockDim.x){
        y[b * D + i] = x[b * D + i] * g;
    }
}

torch::Tensor proto_cos_gate_cuda(torch::Tensor x, torch::Tensor proto, torch::Tensor alpha, torch::Tensor beta){
    auto xc = x.contiguous();
    auto pc = proto.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 b = beta.item<float>();
    int block = 256;
    int grid = B;
    size_t WN = (size_t)((block + 31)>>5);
    size_t shmem = WN * 3 * sizeof(float);
    proto_cos_gate_kernel<<<grid, block, shmem>>>(xc.data_ptr<float>(), pc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, b);
    return y;
}
"""

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

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

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