import torch
from torch.utils.cpp_extension import load_inline

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

// 纯CUDA实现：CosineDistance + Softmax融合
__global__ void compute_cosine_similarities(
    const float* __restrict__ x,
    const float* __restrict__ y,
    const float* __restrict__ x_norms,
    const float* __restrict__ y_norms,
    float* __restrict__ cosine_sims,
    int batch_size,
    int feature_dim,
    float eps
) {
    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;
    
    // 预计算指针
    const float* x_ptr = x + sample_idx * feature_dim;
    const float* y_ptr = y + sample_idx * feature_dim;
    
    // Warp级计算点积
    float dot_product = 0.0f;
    
    // 使用float4向量化
    const float4* x_vec = reinterpret_cast<const float4*>(x_ptr);
    const float4* y_vec = reinterpret_cast<const float4*>(y_ptr);
    
    int feature_dim_vec = feature_dim / 4;
    
    for (int vec_idx = lane_id; vec_idx < feature_dim_vec; vec_idx += 32) {
        float4 x_val = x_vec[vec_idx];
        float4 y_val = y_vec[vec_idx];
        
        dot_product += x_val.x*y_val.x + x_val.y*y_val.y + x_val.z*y_val.z + x_val.w*y_val.w;
    }
    
    // 处理剩余元素
    int remaining_start = feature_dim_vec * 4;
    for (int dim = remaining_start + lane_id; dim < feature_dim; dim += 32) {
        dot_product += x_ptr[dim] * y_ptr[dim];
    }
    
    // Warp级归约
    for (int offset = 16; offset > 0; offset /= 2) {
        dot_product += __shfl_down_sync(0xffffffff, dot_product, offset);
    }
    
    // 使用共享内存进行跨warp归约
    extern __shared__ float shared_data[];
    if (lane_id == 0) {
        shared_data[warp_id] = dot_product;
    }
    __syncthreads();
    
    // 第一个线程计算余弦相似度
    if (tid == 0) {
        float total_dot = 0.0f;
        int num_warps = blockDim.x / 32;
        for (int i = 0; i < num_warps; i++) {
            total_dot += shared_data[i];
        }
        
        float norm_product = x_norms[sample_idx] * y_norms[sample_idx];
        if (norm_product < eps) {
            norm_product = eps;
        }
        
        cosine_sims[sample_idx] = total_dot / norm_product;
    }
}

// 纯CUDA Softmax实现
__global__ void softmax_weighted_distance(
    const float* __restrict__ cosine_sims,
    float* __restrict__ weighted_distances,
    int batch_size,
    float temperature
) {
    int tid = threadIdx.x;
    
    // 使用共享内存进行Softmax计算
    extern __shared__ float shared_softmax[];
    
    // Phase 1: 找到最大值
    float max_val = -FLT_MAX;
    for (int idx = tid; idx < batch_size; idx += blockDim.x) {
        max_val = fmaxf(max_val, cosine_sims[idx] / temperature);
    }
    
    // Warp级归约
    for (int offset = 16; offset > 0; offset /= 2) {
        max_val = fmaxf(max_val, __shfl_down_sync(0xffffffff, max_val, offset));
    }
    
    if (tid % 32 == 0) {
        shared_softmax[tid / 32] = max_val;
    }
    __syncthreads();
    
    // 获取全局最大值
    float global_max = shared_softmax[0];
    for (int i = 1; i < blockDim.x / 32; i++) {
        global_max = fmaxf(global_max, shared_softmax[i]);
    }
    shared_softmax[0] = global_max;
    __syncthreads();
    
    // Phase 2: 计算exp和sum
    float sum_exp = 0.0f;
    for (int idx = tid; idx < batch_size; idx += blockDim.x) {
        float exp_val = expf(cosine_sims[idx] / temperature - global_max);
        sum_exp += exp_val;
        shared_softmax[32 + idx] = exp_val; // 存储exp值
    }
    __syncthreads();
    
    // Warp级归约求和
    for (int offset = 16; offset > 0; offset /= 2) {
        sum_exp += __shfl_down_sync(0xffffffff, sum_exp, offset);
    }
    
    if (tid % 32 == 0) {
        shared_softmax[tid / 32 + 8] = sum_exp;
    }
    __syncthreads();
    
    // 获取全局和
    float global_sum = shared_softmax[8];
    for (int i = 1; i < blockDim.x / 32; i++) {
        global_sum += shared_softmax[i + 8];
    }
    shared_softmax[8] = global_sum;
    __syncthreads();
    
    // Phase 3: 计算加权距离
    for (int idx = tid; idx < batch_size; idx += blockDim.x) {
        float softmax_weight = shared_softmax[32 + idx] / global_sum;
        float cosine_dist = 1.0f - cosine_sims[idx];
        weighted_distances[idx] = softmax_weight * cosine_dist;
    }
}

// 纯CUDA范数计算
__global__ void compute_norms(
    const float* __restrict__ data,
    float* __restrict__ norms,
    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;
    
    const float* data_ptr = data + sample_idx * feature_dim;
    
    // Warp级计算范数平方
    float norm_sq = 0.0f;
    
    const float4* data_vec = reinterpret_cast<const float4*>(data_ptr);
    int feature_dim_vec = feature_dim / 4;
    
    for (int vec_idx = lane_id; vec_idx < feature_dim_vec; vec_idx += 32) {
        float4 val = data_vec[vec_idx];
        norm_sq += val.x*val.x + val.y*val.y + val.z*val.z + val.w*val.w;
    }
    
    // 处理剩余元素
    int remaining_start = feature_dim_vec * 4;
    for (int dim = remaining_start + lane_id; dim < feature_dim; dim += 32) {
        float val = data_ptr[dim];
        norm_sq += val * val;
    }
    
    // Warp级归约
    for (int offset = 16; offset > 0; offset /= 2) {
        norm_sq += __shfl_down_sync(0xffffffff, norm_sq, offset);
    }
    
    // 使用共享内存进行跨warp归约
    extern __shared__ float shared_norms[];
    if (lane_id == 0) {
        shared_norms[warp_id] = norm_sq;
    }
    __syncthreads();
    
    // 第一个线程计算最终范数
    if (tid == 0) {
        float total_norm_sq = 0.0f;
        int num_warps = blockDim.x / 32;
        for (int i = 0; i < num_warps; i++) {
            total_norm_sq += shared_norms[i];
        }
        norms[sample_idx] = sqrtf(total_norm_sq);
    }
}

// 主函数 - 纯CUDA实现
torch::Tensor cosinedistance_softmax_cuda(
    torch::Tensor x,
    torch::Tensor y,
    float eps,
    float temperature
) {
    // 输入验证
    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 cosine_sims = torch::empty({batch_size}, x.options());
    auto weighted_distances = torch::empty({batch_size}, x.options());
    auto x_norms = torch::empty({batch_size}, x.options());
    auto y_norms = torch::empty({batch_size}, x.options());
    
    const int block_size = 256;
    
    // 第一阶段：计算范数
    size_t shared_mem_norms = 8 * sizeof(float);
    compute_norms<<<batch_size, block_size, shared_mem_norms>>>(
        x_contig.data_ptr<float>(),
        x_norms.data_ptr<float>(),
        batch_size,
        feature_dim
    );
    
    compute_norms<<<batch_size, block_size, shared_mem_norms>>>(
        y_contig.data_ptr<float>(),
        y_norms.data_ptr<float>(),
        batch_size,
        feature_dim
    );
    
    // 第二阶段：计算余弦相似度
    size_t shared_mem_cosine = 8 * sizeof(float);
    compute_cosine_similarities<<<batch_size, block_size, shared_mem_cosine>>>(
        x_contig.data_ptr<float>(),
        y_contig.data_ptr<float>(),
        x_norms.data_ptr<float>(),
        y_norms.data_ptr<float>(),
        cosine_sims.data_ptr<float>(),
        batch_size,
        feature_dim,
        eps
    );
    
    // 第三阶段：Softmax加权距离计算
    size_t shared_mem_softmax = (batch_size + 64) * sizeof(float);
    softmax_weighted_distance<<<1, block_size, shared_mem_softmax>>>(
        cosine_sims.data_ptr<float>(),
        weighted_distances.data_ptr<float>(),
        batch_size,
        temperature
    );
    
    return weighted_distances;
}
"""

cosinedistance_softmax_cpp_source = """
torch::Tensor cosinedistance_softmax_cuda(torch::Tensor x, torch::Tensor y, float eps, float temperature);
"""

# 编译CUDA代码
cosinedistance_softmax = load_inline(
    name="cosinedistance_softmax",
    cpp_sources=cosinedistance_softmax_cpp_source,
    cuda_sources=cosinedistance_softmax_source,
    functions=["cosinedistance_softmax_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, eps=1e-8, temperature=1.0):
        super(ModelNew, self).__init__()
        self.eps = eps
        self.temperature = temperature
        self.cosinedistance_softmax = cosinedistance_softmax

    def forward(self, x, y):
        return self.cosinedistance_softmax.cosinedistance_softmax_cuda(x, y, self.eps, self.temperature)
