import torch
from torch.utils.cpp_extension import load_inline

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

// Warp级归约版本 - 融合了Dice Loss和BCE Loss
__global__ void dice_bce_kernel_warp_reduction(
    const float* __restrict__ logits,
    const float* __restrict__ target,
    float* __restrict__ intersection_sum,
    float* __restrict__ pred_sum,
    float* __restrict__ target_sum,
    float* __restrict__ bce_sum,
    int size
) {
    int tid = threadIdx.x;
    int warp_id = (blockIdx.x * blockDim.x + tid) / 32;
    int lane_id = tid % 32;
    
    int elements_per_warp = 1024;
    int warp_start = warp_id * elements_per_warp;
    int warp_end = min(warp_start + elements_per_warp, size);
    
    float reg_intersection = 0.0f;
    float reg_pred_sum = 0.0f;
    float reg_target_sum = 0.0f;
    float reg_bce = 0.0f;

    for (int i = warp_start + lane_id; i < warp_end; i += 32) {
        float l = logits[i];
        float t = target[i];
        
        // --- 关键融合点：计算Sigmoid和BCE ---
        // 使用数值稳定的log-sigmoid公式: log(sigmoid(x)) = -log(1 + exp(-x))
        float p = 1.0f / (1.0f + expf(-l)); // Sigmoid
        float log_sigmoid_val = -logf(1.0f + expf(-l)); // log(sigmoid(l))

        // 累加Dice Loss所需项
        reg_intersection += p * t;
        reg_pred_sum += p;
        reg_target_sum += t;

        // 累加BCE Loss所需项
        // BCE = -[t * log(p) + (1-t) * log(1-p)]
        reg_bce -= (t * log_sigmoid_val + (1.0f - t) * logf(1.0f - p));
    }
    
    // 对所有累加器进行Warp内归约
    #pragma unroll
    for (int offset = 16; offset > 0; offset /= 2) {
        reg_intersection += __shfl_down_sync(0xffffffff, reg_intersection, offset);
        reg_pred_sum += __shfl_down_sync(0xffffffff, reg_pred_sum, offset);
        reg_target_sum += __shfl_down_sync(0xffffffff, reg_target_sum, offset);
        reg_bce += __shfl_down_sync(0xffffffff, reg_bce, offset);
    }
    
    if (lane_id == 0) {
        atomicAdd(intersection_sum, reg_intersection);
        atomicAdd(pred_sum, reg_pred_sum);
        atomicAdd(target_sum, reg_target_sum);
        atomicAdd(bce_sum, reg_bce);
    }
}

torch::Tensor dice_bce_cuda(torch::Tensor logits, torch::Tensor target, float dice_weight, float bce_weight) {
    auto size = logits.numel();
    
    // 创建中间结果张量
    auto intersection = torch::zeros(1, logits.options());
    auto pred_sum_tensor = torch::zeros(1, logits.options());
    auto target_sum_tensor = torch::zeros(1, logits.options());
    auto bce_sum_tensor = torch::zeros(1, logits.options());
    
    const int block_size = 256;
    int num_warps_needed = (size + 1024 - 1) / 1024;
    int num_blocks = (num_warps_needed * 32 + block_size - 1) / block_size;
    
    dice_bce_kernel_warp_reduction<<<num_blocks, block_size>>>(
        logits.data_ptr<float>(), 
        target.data_ptr<float>(),
        intersection.data_ptr<float>(),
        pred_sum_tensor.data_ptr<float>(),
        target_sum_tensor.data_ptr<float>(),
        bce_sum_tensor.data_ptr<float>(),
        size
    );
    
    // 在Host端计算最终损失
    float epsilon = 1e-6f;
    float dice_score = (2.0f * intersection.item<float>()) / (pred_sum_tensor.item<float>() + target_sum_tensor.item<float>() + epsilon);
    float loss_dice = 1.0f - dice_score;
    
    float loss_bce = bce_sum_tensor.item<float>() / size;
    
    float loss_total = dice_weight * loss_dice + bce_weight * loss_bce;
    
    return torch::tensor(loss_total, logits.options());
}
"""

dice_bce_cpp_source = """
torch::Tensor dice_bce_cuda(torch::Tensor logits, torch::Tensor target, float dice_weight, float bce_weight);
"""

dice_bce = load_inline(
    name="dice_bce",
    cpp_sources=dice_bce_cpp_source,
    cuda_sources=dice_bce_source,
    functions=["dice_bce_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, dice_weight=1.0, bce_weight=1.0):
        super(ModelNew, self).__init__()
        self.dice_weight = dice_weight
        self.bce_weight = bce_weight
        self.dice_bce = dice_bce

    def forward(self, logits, target):
        return self.dice_bce.dice_bce_cuda(logits, target, self.dice_weight, self.bce_weight)
