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

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

__global__ void focal_loss_forward_atomic_kernel(
    const float* logits, 
    const float* targets, 
    int batch_size,
    float alpha, 
    float gamma,
    float* final_loss_ptr
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < batch_size) {
        float logit = logits[idx];
        float target = targets[idx];
        
        float prob;
        if (logit > 0) {
            prob = 1.0f / (1.0f + expf(-logit));
        } else {
            float exp_logit = expf(logit);
            prob = exp_logit / (1.0f + exp_logit);
        }
        const float eps = 1e-7f;
        prob = fmaxf(fminf(prob, 1.0f - eps), eps);
        
        float one_minus_prob = 1.0f - prob;
        float log_prob = logf(prob);
        float log_one_minus_prob = logf(one_minus_prob);
        
        float pt = (target > 0.5f) ? prob : one_minus_prob;
        float bce = -(target * log_prob + (1.0f - target) * log_one_minus_prob);
        
        // --- 关键修复：用exp(log(x))重构pow(x, gamma) ---
        float one_minus_pt = 1.0f - pt;
        float focal_weight;
        if (gamma == 2.0f) {
            // 对于gamma=2.0，直接乘法更快更精确
            focal_weight = alpha * one_minus_pt * one_minus_pt;
        } else {
            // 对于其他gamma，使用exp(log)方法
            focal_weight = alpha * expf(gamma * logf(one_minus_pt));
        }
        
        float my_loss = focal_weight * bce;

        atomicAdd(final_loss_ptr, my_loss);
    }
}

torch::Tensor focal_loss_reduction_fused_cuda(
    torch::Tensor logits, 
    torch::Tensor targets, 
    float alpha, 
    float gamma,
    std::string reduction
) {
    logits = logits.contiguous().to(torch::kFloat32);
    targets = targets.contiguous().to(torch::kFloat32);
    
    auto batch_size = logits.numel();
    auto final_loss_tensor = torch::zeros({1}, logits.options());

    const int block_size = 256;
    int num_blocks = (batch_size + block_size - 1) / block_size;

    focal_loss_forward_atomic_kernel<<<num_blocks, block_size>>>(
        logits.data_ptr<float>(),
        targets.data_ptr<float>(),
        batch_size,
        alpha,
        gamma,
        final_loss_tensor.data_ptr<float>()
    );

    if (reduction == "mean") {
        final_loss_tensor = final_loss_tensor / batch_size;
    }

    return final_loss_tensor;
}
"""

focal_loss_fused_cpp_source = """
torch::Tensor focal_loss_reduction_fused_cuda(
    torch::Tensor logits, 
    torch::Tensor targets, 
    float alpha, 
    float gamma,
    std::string reduction
);
"""

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

focal_loss_module = load_inline(
    name="focal_loss_reduction_fused_final",
    cpp_sources=focal_loss_fused_cpp_source,
    cuda_sources=focal_loss_fused_source,
    functions=["focal_loss_reduction_fused_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

    def forward(self, logits, targets):
        logits = logits.to(torch.float32)
        targets = targets.to(torch.float32)

        if logits.dim() == 2 and logits.size(1) == 1:
            logits = logits.squeeze(1)
        
        loss = focal_loss_module.focal_loss_reduction_fused_cuda(
            logits, targets, self.alpha, self.gamma, self.reduction
        )
        
        return loss
