import torch
from torch.utils.cpp_extension import load_inline

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

// 简化高效版本 - 专注核心优化
__global__ void tripletloss_kernel_simple(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    float* __restrict__ losses,
    int batch_size,
    int feature_dim,
    float margin
) {
    // 每个block处理一个样本 - 减少同步
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    int base = sample_idx * feature_dim;
    
    // 计算两个距离 - 4元素展开
    float dist_pos = 0.0f;
    float dist_neg = 0.0f;
    
    int dim = 0;
    for (; dim <= feature_dim - 4; dim += 4) {
        // 同时计算两个距离
        float a0 = anchor[base + dim];
        float a1 = anchor[base + dim + 1];
        float a2 = anchor[base + dim + 2];
        float a3 = anchor[base + dim + 3];
        
        float p0 = positive[base + dim];
        float p1 = positive[base + dim + 1];
        float p2 = positive[base + dim + 2];
        float p3 = positive[base + dim + 3];
        
        float n0 = negative[base + dim];
        float n1 = negative[base + dim + 1];
        float n2 = negative[base + dim + 2];
        float n3 = negative[base + dim + 3];
        
        // 计算差值并累加
        float diff_pos0 = a0 - p0;
        float diff_pos1 = a1 - p1;
        float diff_pos2 = a2 - p2;
        float diff_pos3 = a3 - p3;
        
        float diff_neg0 = a0 - n0;
        float diff_neg1 = a1 - n1;
        float diff_neg2 = a2 - n2;
        float diff_neg3 = a3 - n3;
        
        dist_pos += diff_pos0*diff_pos0 + diff_pos1*diff_pos1 + diff_pos2*diff_pos2 + diff_pos3*diff_pos3;
        dist_neg += diff_neg0*diff_neg0 + diff_neg1*diff_neg1 + diff_neg2*diff_neg2 + diff_neg3*diff_neg3;
    }
    
    // 处理剩余元素
    for (; dim < feature_dim; dim++) {
        float a_val = anchor[base + dim];
        float p_val = positive[base + dim];
        float n_val = negative[base + dim];
        
        float diff_pos = a_val - p_val;
        float diff_neg = a_val - n_val;
        
        dist_pos += diff_pos * diff_pos;
        dist_neg += diff_neg * diff_neg;
    }
    
    // 计算最终距离和loss
    dist_pos = sqrtf(dist_pos);
    dist_neg = sqrtf(dist_neg);
    
    float sample_loss = dist_pos - dist_neg + margin;
    if (sample_loss > 0.0f) {
        losses[sample_idx] = sample_loss;
    } else {
        losses[sample_idx] = 0.0f;
    }
}

// 向量化版本 - 使用float4
__global__ void tripletloss_kernel_vectorized(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    float* __restrict__ losses,
    int batch_size,
    int feature_dim,
    float margin
) {
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    // 使用float4向量化
    const float4* anchor_vec = reinterpret_cast<const float4*>(anchor);
    const float4* positive_vec = reinterpret_cast<const float4*>(positive);
    const float4* negative_vec = reinterpret_cast<const float4*>(negative);
    
    int feature_dim_vec = feature_dim / 4;
    int base_vec = sample_idx * feature_dim_vec;
    
    float dist_pos = 0.0f;
    float dist_neg = 0.0f;
    
    // 向量化计算
    for (int dim_vec = 0; dim_vec < feature_dim_vec; dim_vec++) {
        float4 a_vec = anchor_vec[base_vec + dim_vec];
        float4 p_vec = positive_vec[base_vec + dim_vec];
        float4 n_vec = negative_vec[base_vec + dim_vec];
        
        float diff_pos_x = a_vec.x - p_vec.x;
        float diff_pos_y = a_vec.y - p_vec.y;
        float diff_pos_z = a_vec.z - p_vec.z;
        float diff_pos_w = a_vec.w - p_vec.w;
        
        float diff_neg_x = a_vec.x - n_vec.x;
        float diff_neg_y = a_vec.y - n_vec.y;
        float diff_neg_z = a_vec.z - n_vec.z;
        float diff_neg_w = a_vec.w - n_vec.w;
        
        dist_pos += diff_pos_x*diff_pos_x + diff_pos_y*diff_pos_y + diff_pos_z*diff_pos_z + diff_pos_w*diff_pos_w;
        dist_neg += diff_neg_x*diff_neg_x + diff_neg_y*diff_neg_y + diff_neg_z*diff_neg_z + diff_neg_w*diff_neg_w;
    }
    
    // 处理剩余元素
    for (int dim = feature_dim_vec * 4; dim < feature_dim; dim++) {
        float a_val = anchor[sample_idx * feature_dim + dim];
        float p_val = positive[sample_idx * feature_dim + dim];
        float n_val = negative[sample_idx * feature_dim + dim];
        
        float diff_pos = a_val - p_val;
        float diff_neg = a_val - n_val;
        
        dist_pos += diff_pos * diff_pos;
        dist_neg += diff_neg * diff_neg;
    }
    
    // 计算最终距离和loss
    dist_pos = sqrtf(dist_pos);
    dist_neg = sqrtf(dist_neg);
    
    float sample_loss = dist_pos - dist_neg + margin;
    if (sample_loss > 0.0f) {
        losses[sample_idx] = sample_loss;
    } else {
        losses[sample_idx] = 0.0f;
    }
}

torch::Tensor tripletloss_cuda(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    float margin,
    std::string mode = "simple"
) {
    TORCH_CHECK(anchor.scalar_type() == torch::kFloat32, "Anchor must be float32");
    TORCH_CHECK(positive.scalar_type() == torch::kFloat32, "Positive must be float32");
    TORCH_CHECK(negative.scalar_type() == torch::kFloat32, "Negative must be float32");
    
    auto anchor_contig = anchor.contiguous();
    auto positive_contig = positive.contiguous();
    auto negative_contig = negative.contiguous();
    
    int batch_size = anchor_contig.size(0);
    int feature_dim = anchor_contig.size(1);
    
    // 直接输出每个样本的loss
    auto losses = torch::zeros({batch_size}, anchor.options());
    
    if (mode == "vectorized") {
        // 向量化版本
        tripletloss_kernel_vectorized<<<batch_size, 1>>>(
            anchor_contig.data_ptr<float>(),
            positive_contig.data_ptr<float>(),
            negative_contig.data_ptr<float>(),
            losses.data_ptr<float>(),
            batch_size,
            feature_dim,
            margin
        );
    } else {
        // 简化版本（默认）
        tripletloss_kernel_simple<<<batch_size, 1>>>(
            anchor_contig.data_ptr<float>(),
            positive_contig.data_ptr<float>(),
            negative_contig.data_ptr<float>(),
            losses.data_ptr<float>(),
            batch_size,
            feature_dim,
            margin
        );
    }
    
    
    return torch::mean(losses);
}
"""

tripletloss_cpp_source = """
torch::Tensor tripletloss_cuda(torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, float margin, std::string mode);
"""

# 简化编译选项
tripletloss = load_inline(
    name="tripletloss",
    cpp_sources=tripletloss_cpp_source,
    cuda_sources=tripletloss_source,
    functions=["tripletloss_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, margin=1.0, p=2.0, eps=1e-6, mode="vectorized"):
        super(ModelNew, self).__init__()
        self.margin = margin
        self.p = p
        self.eps = eps
        self.mode = mode
        self.tripletloss = tripletloss

    def forward(self, anchor, positive, negative):
        return self.tripletloss.tripletloss_cuda(
            anchor, positive, negative, self.margin, self.mode
        )
