import torch
from torch.utils.cpp_extension import load_inline

# CUDA 源代码字符串
cosine_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

// 融合内核：余弦相似度 + 对比损失
__global__ void cosine_contrastive_loss_kernel(
    const float* __restrict__ x,
    const float* __restrict__ y,
    const float* __restrict__ labels, // 新增：标签
    float* __restrict__ losses,       // 输出：每个样本的损失
    int batch_size,
    int feature_dim,
    float margin                      // 新增：margin值
) {
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    int tid = threadIdx.x;
    int base = sample_idx * feature_dim;
    
    // --- 1. 采样计算余弦相似度 (复用你之前的优秀代码) ---
    const float SAMPLE_RATIO = 0.7f;
    const int stride = max(1, (int)(1.0f / SAMPLE_RATIO));
    
    float sampled_dot = 0.0f;
    float sampled_norm_x = 0.0f;
    float sampled_norm_y = 0.0f;
    
    for (int i = tid; i < feature_dim; i += blockDim.x * stride) {
        float x_val = x[base + i];
        float y_val = y[base + i];
        
        sampled_dot += x_val * y_val;
        sampled_norm_x += x_val * x_val;
        sampled_norm_y += y_val * y_val;
    }
    
    for (int offset = 16; offset > 0; offset /= 2) {
        sampled_dot += __shfl_down_sync(0xffffffff, sampled_dot, offset);
        sampled_norm_x += __shfl_down_sync(0xffffffff, sampled_norm_x, offset);
        sampled_norm_y += __shfl_down_sync(0xffffffff, sampled_norm_y, offset);
    }

    float cosine_sim = 0.0f;
    if (tid == 0) {
        float scale_factor = 1.0f / SAMPLE_RATIO;
        float total_dot = sampled_dot * scale_factor;
        float total_norm_x = sampled_norm_x * scale_factor;
        float total_norm_y = sampled_norm_y * scale_factor;
        
        float norm_x = sqrtf(total_norm_x);
        float norm_y = sqrtf(total_norm_y);
        cosine_sim = total_dot / (norm_x * norm_y + 1e-8f);

        // --- 2. 融合：直接计算对比损失 ---
        float label = labels[sample_idx];
        float loss_val;
        if (label > 0.5f) { // 假设标签为1表示相似
            loss_val = 1.0f - cosine_sim;
        } else {
            loss_val = fmaxf(0.0f, cosine_sim - margin);
        }
        losses[sample_idx] = loss_val;
    }
}

torch::Tensor cosine_contrastive_loss_cuda(
    torch::Tensor x,
    torch::Tensor y,
    torch::Tensor labels,
    float margin
) {
    auto x_contig = x.contiguous();
    auto y_contig = y.contiguous();
    auto labels_contig = labels.contiguous();
    
    int batch_size = x_contig.size(0);
    int feature_dim = x_contig.size(1);
    
    // 输出每个样本的损失，最后在Python中求和
    auto losses = torch::zeros({batch_size}, x.options());
    
    const int block_size = 32;
    
    cosine_contrastive_loss_kernel<<<batch_size, block_size>>>(
        x_contig.data_ptr<float>(),
        y_contig.data_ptr<float>(),
        labels_contig.data_ptr<float>(),
        losses.data_ptr<float>(),
        batch_size,
        feature_dim,
        margin
    );
    
    return losses;
}
"""

cosine_cpp_source = """
torch::Tensor cosine_contrastive_loss_cuda(torch::Tensor x, torch::Tensor y, torch::Tensor labels, float margin);
"""

# 编译CUDA代码
cosine = load_inline(
    name="cosine",
    cpp_sources=cosine_cpp_source,
    cuda_sources=cosine_source,
    functions=["cosine_contrastive_loss_cuda"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-std=c++17"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, margin=0.5):
        super(ModelNew, self).__init__()
        self.cosine = cosine
        self.margin = margin

    def forward(self, x, y, labels):
        # 内核返回每个样本的loss，我们手动求和
        per_sample_loss = self.cosine.cosine_contrastive_loss_cuda(x, y, labels, self.margin)
        return torch.sum(per_sample_loss)
