import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

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

__global__ void cosine_swish_gelu_kernel(
    const float* __restrict__ x, 
    const float* __restrict__ target, 
    float* __restrict__ y, 
    int batch_size, 
    int width) 
{
    int row = blockIdx.x;
    int tid = threadIdx.x;

    if (row >= batch_size) return;

    const float* row_x = x + row * width;

    float sum_dot = 0.0f;
    float sum_xx = 0.0f;
    float sum_tt = 0.0f;

    for (int i = tid; i < width; i += blockDim.x) {
        float val_x = row_x[i];
        float val_t = target[i];

        sum_dot += val_x * val_t;
        sum_xx += val_x * val_x;
        sum_tt += val_t * val_t;
    }

    sum_dot = warp_reduce(sum_dot);
    sum_xx = warp_reduce(sum_xx);
    sum_tt = warp_reduce(sum_tt);

    static __shared__ float shared_dot[32];
    static __shared__ float shared_xx[32];
    static __shared__ float shared_tt[32];

    int lane = tid % 32;
    int wid = tid / 32;

    if (lane == 0) {
        shared_dot[wid] = sum_dot;
        shared_xx[wid] = sum_xx;
        shared_tt[wid] = sum_tt;
    }
    __syncthreads();

    sum_dot = (tid < blockDim.x / 32) ? shared_dot[lane] : 0.0f;
    sum_xx = (tid < blockDim.x / 32) ? shared_xx[lane] : 0.0f;
    sum_tt = (tid < blockDim.x / 32) ? shared_tt[lane] : 0.0f;

    if (wid == 0) {
        sum_dot = warp_reduce(sum_dot);
        sum_xx = warp_reduce(sum_xx);
        sum_tt = warp_reduce(sum_tt);
    }

    if (tid == 0) {
        float norm_x = sqrtf(sum_xx);
        float norm_t = sqrtf(sum_tt);
        float cosine = sum_dot / (norm_x * norm_t + 1e-6f);

        float swish = cosine / (1.0f + expf(-cosine));
        float gelu = swish * 0.5f * (1.0f + erff(swish * 0.70710678f));

        y[row] = gelu;
    }
}

torch::Tensor launch_cosine_swish_gelu(torch::Tensor x, torch::Tensor target) {
    auto batch_size = x.size(0);
    auto width = x.size(1);
    auto y = torch::empty({batch_size}, x.options());

    const int threads = 256;
    const int blocks = batch_size;

    cosine_swish_gelu_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(), 
        target.data_ptr<float>(), 
        y.data_ptr<float>(), 
        batch_size, 
        width
    );
    return y;
}
"""

cpp_source = """
torch::Tensor launch_cosine_swish_gelu(torch::Tensor x, torch::Tensor target);
"""

cosine_swish_gelu_module = load_inline(
    name='cosine_swish_gelu_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_cosine_swish_gelu'],
    verbose=False
)


class ModelNew(nn.Module):
    def __init__(self, target):
        super(ModelNew, self).__init__()
        self.target = nn.Parameter(target)
        self.op = cosine_swish_gelu_module

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.op.launch_cosine_swish_gelu(x.contiguous(), self.target.contiguous())