import torch
from torch.utils.cpp_extension import load_inline

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

// 智能优化版本 - 动态调整block数量减少原子操作竞争
__global__ void contrastiveloss_kernel_smart(
    const float* __restrict__ anchor,
    const float* __restrict__ sample,
    const float* __restrict__ label,
    float* __restrict__ contrastive_loss,
    int batch_size,
    int feature_dim,
    float margin
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < batch_size) {
        // 计算欧氏距离的平方
        float distance_sq = 0.0f;
        int anchor_offset = idx * feature_dim;
        int sample_offset = idx * feature_dim;
        
        // 融合计算：减法 + 平方 + 求和
        for (int i = 0; i < feature_dim; i++) {
            float diff = anchor[anchor_offset + i] - sample[sample_offset + i];
            distance_sq += diff * diff;
        }
        
        // 计算实际欧氏距离
        float distance = sqrtf(distance_sq);
        
        // 标准ContrastiveLoss公式
        float label_val = label[idx];
        float positive_loss = 0.5f * distance_sq * label_val;
        float margin_diff = margin - distance;
        float clamped_diff = fmaxf(0.0f, margin_diff);
        float negative_loss = 0.5f * clamped_diff * clamped_diff * (1.0f - label_val);
        
        float loss = positive_loss + negative_loss;
        
        // 原子操作累加到总损失
        atomicAdd(contrastive_loss, loss);
    }
}

// 高效归约版本 - 解决原子操作瓶颈
__global__ void contrastiveloss_kernel_efficient(
    const float* __restrict__ anchor,
    const float* __restrict__ sample,
    const float* __restrict__ label,
    float* __restrict__ partial_sums,
    int batch_size,
    int feature_dim,
    float margin
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 局部累加器
    float local_loss = 0.0f;
    
    // 每个线程处理多个样本
    int stride = blockDim.x * gridDim.x;
    for (int sample_idx = idx; sample_idx < batch_size; sample_idx += stride) {
        // 计算欧氏距离的平方
        float distance_sq = 0.0f;
        int anchor_offset = sample_idx * feature_dim;
        int sample_offset = sample_idx * feature_dim;
        
        for (int i = 0; i < feature_dim; i++) {
            float diff = anchor[anchor_offset + i] - sample[sample_offset + i];
            distance_sq += diff * diff;
        }
        
        // 计算实际欧氏距离
        float distance = sqrtf(distance_sq);
        
        // 标准ContrastiveLoss公式
        float label_val = label[sample_idx];
        float positive_loss = 0.5f * distance_sq * label_val;
        float margin_diff = margin - distance;
        float clamped_diff = fmaxf(0.0f, margin_diff);
        float negative_loss = 0.5f * clamped_diff * clamped_diff * (1.0f - label_val);
        
        local_loss += positive_loss + negative_loss;
    }
    
    // 使用共享内存进行块内归约
    extern __shared__ float shared_mem[];
    shared_mem[threadIdx.x] = local_loss;
    
    __syncthreads();
    
    // 块内归约
    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
        if (threadIdx.x < stride) {
            shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
        }
        __syncthreads();
    }
    
    // 每个block写入部分和到全局内存
    if (threadIdx.x == 0) {
        partial_sums[blockIdx.x] = shared_mem[0];
    }
}

torch::Tensor contrastiveloss_cuda(
    torch::Tensor anchor, 
    torch::Tensor sample, 
    torch::Tensor label, 
    float margin, 
    std::string mode = "smart"
) {
    TORCH_CHECK(anchor.scalar_type() == torch::kFloat32, "Anchor must be float32");
    TORCH_CHECK(sample.scalar_type() == torch::kFloat32, "Sample must be float32");
    TORCH_CHECK(label.scalar_type() == torch::kFloat32, "Label must be float32");
    TORCH_CHECK(anchor.sizes() == sample.sizes(), "Anchor and sample must have same shape");
    TORCH_CHECK(anchor.size(0) == label.size(0), "Batch size mismatch");
    
    auto anchor_contig = anchor.contiguous();
    auto sample_contig = sample.contiguous();
    auto label_contig = label.contiguous();
    
    int batch_size = anchor_contig.size(0);
    int feature_dim = anchor_contig.size(1);
    
    if (mode == "efficient") {
        // 高效归约版本 - 解决原子操作瓶颈
        const int block_size = 256;
        
        // 智能选择block数量：确保每个block处理适量样本
        int num_blocks;
        if (batch_size <= 256) {
            num_blocks = 4;  // 关键：强制使用多个blocks
        } else if (batch_size <= 512) {
            num_blocks = 8;
        } else if (batch_size <= 1024) {
            num_blocks = 16;
        } else {
            num_blocks = min(32, (batch_size + 255) / 256);
        }
        
        // 创建部分和数组
        auto partial_sums = torch::zeros({num_blocks}, anchor.options());
        
        size_t shared_mem = block_size * sizeof(float);
        contrastiveloss_kernel_efficient<<<num_blocks, block_size, shared_mem>>>(
            anchor_contig.data_ptr<float>(),
            sample_contig.data_ptr<float>(),
            label_contig.data_ptr<float>(),
            partial_sums.data_ptr<float>(),
            batch_size,
            feature_dim,
            margin
        );
        
        // 在GPU上完成最终归约（避免CPU-GPU传输）
        auto total_loss = torch::sum(partial_sums);
        return total_loss;
        
    } else {
        // 智能优化版本 - 动态调整block数量
        auto contrastive_loss = torch::zeros(1, anchor.options());
        
        const int block_size = 256;
        
        // 智能block分配策略
        int num_blocks;
        if (batch_size <= 64) {
            num_blocks = 1;
        } else if (batch_size <= 128) {
            num_blocks = 1;
        } else if (batch_size <= 256) {
            num_blocks = 2;  // 关键：增加block数量减少竞争
        } else if (batch_size <= 512) {
            num_blocks = 4;
        } else {
            num_blocks = min(16, (batch_size + 255) / 256);
        }
        
        contrastiveloss_kernel_smart<<<num_blocks, block_size>>>(
            anchor_contig.data_ptr<float>(),
            sample_contig.data_ptr<float>(),
            label_contig.data_ptr<float>(),
            contrastive_loss.data_ptr<float>(),
            batch_size,
            feature_dim,
            margin
        );
        
        return contrastive_loss;
    }
}
"""

contrastiveloss_cpp_source = """
torch::Tensor contrastiveloss_cuda(torch::Tensor anchor, torch::Tensor sample, torch::Tensor label, float margin, std::string mode);
"""

# 编译CUDA代码
contrastiveloss = load_inline(
    name="contrastiveloss",
    cpp_sources=contrastiveloss_cpp_source,
    cuda_sources=contrastiveloss_source,
    functions=["contrastiveloss_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=2.0, mode="smart"):
        super(ModelNew, self).__init__()
        self.margin = margin
        self.mode = mode
        self.contrastiveloss = contrastiveloss

    def forward(self, anchor, sample, label):
        return self.contrastiveloss.contrastiveloss_cuda(anchor, sample, label, self.margin, self.mode)
