import torch
from torch.utils.cpp_extension import load_inline
import os

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

__global__ void focal_loss_kernel(
    const float* logits, 
    const float* targets, 
    float* loss, 
    int batch_size,
    float alpha, 
    float gamma
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < batch_size) {
        float logit = logits[idx];
        float target = targets[idx];  // 改为float类型
        
        // Compute probability using sigmoid with numerical stability
        float prob;
        if (logit > 0) {
            float exp_neg_logit = expf(-logit);
            prob = 1.0f / (1.0f + exp_neg_logit);
        } else {
            float exp_logit = expf(logit);
            prob = exp_logit / (1.0f + exp_logit);
        }
        
        // Clamp probability to avoid log(0)
        prob = fmaxf(fminf(prob, 1.0f - 1e-7f), 1e-7f);
        
        // Compute pt
        float pt = (target > 0.5f) ? prob : (1.0f - prob);  // 改为float比较
        
        // Compute binary cross entropy
        float bce = -(target * logf(prob) + (1.0f - target) * logf(1.0f - prob));
        
        // Compute focal weight
        float focal_weight = alpha * powf(1.0f - pt, gamma);
        
        // Apply focal weight
        loss[idx] = focal_weight * bce;
    }
}

torch::Tensor focal_loss_cuda(torch::Tensor logits, torch::Tensor targets, float alpha, float gamma) {
    // 确保输入输出数据类型完全一致
    logits = logits.to(torch::kFloat32);
    targets = targets.to(torch::kFloat32);  // 改为float32
    
    auto batch_size = logits.numel();
    auto loss = torch::empty_like(logits);
    
    // 动态计算线程块数量
    const int block_size = 256;
    int num_blocks = (batch_size + block_size - 1) / block_size;
    
    // 启动核函数
    focal_loss_kernel<<<num_blocks, block_size>>>(
        logits.data_ptr<float>(),
        targets.data_ptr<float>(),  // 改为float*
        loss.data_ptr<float>(),
        batch_size,
        alpha,
        gamma
    );
    
    // 检查CUDA错误
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        AT_ERROR("CUDA kernel failed: ", cudaGetErrorString(err));
    }
    
    return loss;
}
"""

focal_loss_cpp_source = """
torch::Tensor focal_loss_cuda(torch::Tensor logits, torch::Tensor targets, float alpha, float gamma);
"""


build_dir = './cuda_build'
os.makedirs(build_dir, exist_ok=True)

# Compile the inline CUDA code
focal_loss_module = load_inline(
    name="focal_loss",
    cpp_sources=focal_loss_cpp_source,
    cuda_sources=focal_loss_source,
    functions=["focal_loss_cuda"],
    verbose=True,
    build_directory=build_dir,
    extra_cuda_cflags=["-O3"] 
)

class ModelNew(torch.nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
        super(ModelNew, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction
        self.focal_loss = focal_loss_module

    def forward(self, logits, targets):
        # 确保输入输出类型完全一致
        logits = logits.to(torch.float32)
        targets = targets.to(torch.float32)  # 改为float32
        
        # 确保logits和targets形状一致
        if logits.dim() == 2 and logits.size(1) == 1:
            logits = logits.squeeze(1)  # 将logits从[batch_size, 1]变为[batch_size]
        
        # Compute focal loss using custom CUDA kernel
        loss = self.focal_loss.focal_loss_cuda(logits, targets, self.alpha, self.gamma)
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:
            return loss
