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

# 定义CUDA内核和C++绑定代码
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

// SELU-Clip前向传播CUDA内核
__global__ void selu_clip_forward_kernel(
    const float* __restrict__ x,
    float* __restrict__ output,
    float alpha,
    float scale,
    float clip_min,
    float clip_max,
    int num_elements) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < num_elements) {
        float val = x[idx];
        // 计算SELU
        float selu_val;
        if (val > 0.0f) {
            selu_val = scale * val;
        } else {
            selu_val = scale * alpha * (expf(val) - 1.0f);
        }
        // 裁剪操作
        if (selu_val < clip_min) {
            output[idx] = clip_min;
        } else if (selu_val > clip_max) {
            output[idx] = clip_max;
        } else {
            output[idx] = selu_val;
        }
    }
}

// SELU-Clip反向传播CUDA内核
__global__ void selu_clip_backward_kernel(
    const float* __restrict__ x,
    const float* __restrict__ grad_output,
    float* __restrict__ grad_input,
    float alpha,
    float scale,
    int num_elements) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < num_elements) {
        float val = x[idx];
        float grad = grad_output[idx];
        // 计算梯度
        if (val > 0.0f) {
            grad_input[idx] = scale * grad;
        } else {
            grad_input[idx] = scale * alpha * expf(val) * grad;
        }
    }
}

// C++绑定前向函数
torch::Tensor selu_clip_forward(
    torch::Tensor x,
    float alpha,
    float scale,
    float clip_min,
    float clip_max) {
    auto output = torch::empty_like(x);
    int num_elements = x.numel();
    int block_size = 256;
    int grid_size = (num_elements + block_size - 1) / block_size;
    
    selu_clip_forward_kernel<<<grid_size, block_size>>>(
        x.data_ptr<float>(),
        output.data_ptr<float>(),
        alpha,
        scale,
        clip_min,
        clip_max,
        num_elements
    );
    return output;
}

// C++绑定反向函数
torch::Tensor selu_clip_backward(
    torch::Tensor x,
    torch::Tensor grad_output,
    float alpha,
    float scale) {
    auto grad_input = torch::empty_like(x);
    int num_elements = x.numel();
    int block_size = 256;
    int grid_size = (num_elements + block_size - 1) / block_size;
    
    selu_clip_backward_kernel<<<grid_size, block_size>>>(
        x.data_ptr<float>(),
        grad_output.data_ptr<float>(),
        grad_input.data_ptr<float>(),
        alpha,
        scale,
        num_elements
    );
    return grad_input;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &selu_clip_forward, "SELU-Clip forward");
    m.def("backward", &selu_clip_backward, "SELU-Clip backward");
}
"""

# 动态编译CUDA内核
selu_clip_cuda = load_inline(
    name="selu_clip_cuda",
    cpp_sources=[],
    cuda_sources=[cuda_source],
    extra_cuda_cflags=["-O2"],
    with_cuda=True,
    verbose=False
)

class SELUClipFunction(Function):
    """自定义CUDA实现的SELU-Clip算子函数"""
    @staticmethod
    def forward(ctx, x, alpha, scale, clip_min, clip_max):
        ctx.alpha = alpha
        ctx.scale = scale
        ctx.save_for_backward(x)
        return selu_clip_cuda.forward(x, alpha, scale, clip_min, clip_max)

    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        alpha = ctx.alpha
        scale = ctx.scale
        grad_input = selu_clip_cuda.backward(x, grad_output, alpha, scale)
        return grad_input, None, None, None, None

class ModelNew(nn.Module):
    """使用内置CUDA内核的SELU-Clip模型"""
    def __init__(self, alpha=1.6732632423543772, scale=1.0507009873554804, clip_min=-5.0, clip_max=5.0):
        super(ModelNew, self).__init__()
        self.alpha = alpha
        self.scale = scale
        self.clip_min = clip_min
        self.clip_max = clip_max

    def forward(self, x):
        return SELUClipFunction.apply(x, self.alpha, self.scale, self.clip_min, self.clip_max)