import torch
from torch.utils.cpp_extension import load_inline

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

// 纯CUDA实现：Chebyshev + Sigmoid融合
__global__ void chebyshev_sigmoid_kernel(
    const float* __restrict__ x,
    const float* __restrict__ y,
    float* __restrict__ distances,
    float* __restrict__ sigmoid_weights,
    float* __restrict__ weighted_distances,
    int batch_size,
    int feature_dim
) {
    // 每个block处理一个样本
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    int tid = threadIdx.x;
    int base = sample_idx * feature_dim;
    
    // 使用共享内存进行归约
    extern __shared__ float shared_max[];
    shared_max[tid] = 0.0f;
    
    // 每个线程处理多个元素
    int stride = blockDim.x;
    for (int dim = tid; dim < feature_dim; dim += stride) {
        float diff = x[base + dim] - y[base + dim];
        float abs_diff = fabsf(diff);
        if (abs_diff > shared_max[tid]) {
            shared_max[tid] = abs_diff;
        }
    }
    
    __syncthreads();
    
    // 块内归约找最大值
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            if (shared_max[tid + stride] > shared_max[tid]) {
                shared_max[tid] = shared_max[tid + stride];
            }
        }
        __syncthreads();
    }
    
    // 第一个线程写入结果
    if (tid == 0) {
        float chebyshev_dist = shared_max[0];
        distances[sample_idx] = chebyshev_dist;
        
        // 应用Sigmoid激活
        float sigmoid_weight = 1.0f / (1.0f + expf(-chebyshev_dist));
        sigmoid_weights[sample_idx] = sigmoid_weight;
        
        // 计算加权距离
        weighted_distances[sample_idx] = sigmoid_weight * chebyshev_dist;
    }
}

// Warp级优化版本 - Chebyshev + Sigmoid融合
__global__ void chebyshev_sigmoid_kernel_warp(
    const float* __restrict__ x,
    const float* __restrict__ y,
    float* __restrict__ distances,
    float* __restrict__ sigmoid_weights,
    float* __restrict__ weighted_distances,
    int batch_size,
    int feature_dim
) {
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    int tid = threadIdx.x;
    int warp_id = tid / 32;
    int lane_id = tid % 32;
    
    int base = sample_idx * feature_dim;
    
    // Warp级计算最大值
    float warp_max = 0.0f;
    
    // 每个warp处理一部分特征
    int elements_per_warp = (feature_dim + 8 - 1) / 8;
    int start_dim = warp_id * elements_per_warp;
    int end_dim = min(start_dim + elements_per_warp, feature_dim);
    
    for (int dim = start_dim + lane_id; dim < end_dim; dim += 32) {
        float diff = x[base + dim] - y[base + dim];
        float abs_diff = fabsf(diff);
        if (abs_diff > warp_max) {
            warp_max = abs_diff;
        }
    }
    
    // Warp级归约找最大值
    for (int offset = 16; offset > 0; offset /= 2) {
        float other_max = __shfl_down_sync(0xffffffff, warp_max, offset);
        if (other_max > warp_max) {
            warp_max = other_max;
        }
    }
    
    // 使用共享内存进行跨warp归约
    extern __shared__ float shared_data[];
    if (lane_id == 0) {
        shared_data[warp_id] = warp_max;
    }
    __syncthreads();
    
    // 第一个线程找到全局最大值并应用Sigmoid
    if (tid == 0) {
        float global_max = 0.0f;
        int num_warps = blockDim.x / 32;
        for (int i = 0; i < num_warps; i++) {
            if (shared_data[i] > global_max) {
                global_max = shared_data[i];
            }
        }
        
        float chebyshev_dist = global_max;
        distances[sample_idx] = chebyshev_dist;
        
        // 应用Sigmoid激活：sigmoid(x) = 1 / (1 + exp(-x))
        float sigmoid_weight = 1.0f / (1.0f + expf(-chebyshev_dist));
        sigmoid_weights[sample_idx] = sigmoid_weight;
        
        // 计算加权距离
        weighted_distances[sample_idx] = sigmoid_weight * chebyshev_dist;
    }
}

// 全局Sigmoid加权版本 - 批量处理
__global__ void sigmoid_weighted_distance_global(
    const float* __restrict__ distances,
    float* __restrict__ sigmoid_weights,
    float* __restrict__ weighted_distances,
    int batch_size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= batch_size) return;
    
    float chebyshev_dist = distances[idx];
    
    // 应用Sigmoid激活
    float sigmoid_weight = 1.0f / (1.0f + expf(-chebyshev_dist));
    sigmoid_weights[idx] = sigmoid_weight;
    
    // 计算加权距离
    weighted_distances[idx] = sigmoid_weight * chebyshev_dist;
}

// 主函数 - 纯CUDA实现
torch::Tensor chebyshev_sigmoid_cuda(
    torch::Tensor x,
    torch::Tensor y
) {
    // 输入验证
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "X must be float32");
    TORCH_CHECK(y.scalar_type() == torch::kFloat32, "Y must be float32");
    TORCH_CHECK(x.sizes() == y.sizes(), "X and Y must have same shape");
    
    auto x_contig = x.contiguous();
    auto y_contig = y.contiguous();
    
    int batch_size = x_contig.size(0);
    int feature_dim = x_contig.size(1);
    
    // 创建输出张量
    auto distances = torch::zeros({batch_size}, x.options());
    auto sigmoid_weights = torch::zeros({batch_size}, x.options());
    auto weighted_distances = torch::zeros({batch_size}, x.options());
    
    const int block_size = 256; // 8个warps
    size_t shared_mem = 8 * sizeof(float); // 8个warp的结果
    
    // 第一阶段：计算Chebyshev距离
    chebyshev_sigmoid_kernel_warp<<<batch_size, block_size, shared_mem>>>(
        x_contig.data_ptr<float>(),
        y_contig.data_ptr<float>(),
        distances.data_ptr<float>(),
        sigmoid_weights.data_ptr<float>(),
        weighted_distances.data_ptr<float>(),
        batch_size,
        feature_dim
    );
    
    return weighted_distances;
}
"""

chebyshev_sigmoid_cpp_source = """
torch::Tensor chebyshev_sigmoid_cuda(torch::Tensor x, torch::Tensor y);
"""

# 编译CUDA代码
chebyshev_sigmoid = load_inline(
    name="chebyshev_sigmoid",
    cpp_sources=chebyshev_sigmoid_cpp_source,
    cuda_sources=chebyshev_sigmoid_source,
    functions=["chebyshev_sigmoid_cuda"],
    extra_cuda_cflags=[
        "-O3", 
        "--use_fast_math",
        "-gencode=arch=compute_80,code=sm_80"
    ],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.chebyshev_sigmoid = chebyshev_sigmoid

    def forward(self, x, y):
        return self.chebyshev_sigmoid.chebyshev_sigmoid_cuda(x, y)
