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

# CUDA C++ 源代码字符串
focal_loss_cpp_source = """
torch::Tensor focal_loss_forward_backward_cuda(
    torch::Tensor logits, 
    torch::Tensor targets, 
    float alpha, 
    float gamma
);
"""

# CUDA 源代码字符串
focal_loss_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>

__global__ void focal_loss_forward_backward_kernel(
    const float* logits, 
    const float* targets, 
    float* loss,          // 输出：损失值
    float* d_logits,       // 输出：梯度
    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];
        
        // --- 1. 前向计算部分 ---
        float prob; // p
        if (logit > 0) {
            prob = 1.0f / (1.0f + expf(-logit));
        } else {
            float exp_logit = expf(logit);
            prob = exp_logit / (1.0f + exp_logit);
        }
        prob = fmaxf(fminf(prob, 1.0f - 1e-7f), 1e-7f);
        
        float pt = (target > 0.5f) ? prob : (1.0f - prob);
        float bce = -(target * logf(prob) + (1.0f - target) * logf(1.0f - prob));
        float focal_weight = alpha * powf(1.0f - pt, gamma);
        loss[idx] = focal_weight * bce;

        // --- 2. 反向梯度计算部分 ---
        // dL/dlogit = alpha * gamma * (1-pt)^(gamma-1) * (1-p) * p * log(pt) - alpha * (1-pt)^gamma * (p - target)
        float one_minus_pt = 1.0f - pt;
        float one_minus_pt_pow_gamma_minus_1 = powf(one_minus_pt, gamma - 1.0f);
        
        float term1 = alpha * gamma * one_minus_pt_pow_gamma_minus_1 * (1.0f - prob) * prob * logf(pt);
        float term2 = focal_weight * (prob - target);
        
        d_logits[idx] = term1 - term2;
    }
}

// C++ 封装函数
torch::Tensor focal_loss_forward_backward_cuda(
    torch::Tensor logits, 
    torch::Tensor targets, 
    float alpha, 
    float gamma
) {
    logits = logits.contiguous().to(torch::kFloat32);
    targets = targets.contiguous().to(torch::kFloat32);
    
    auto batch_size = logits.numel();
    
    // 创建一个输出张量，前半部分存loss，后半部分存梯度
    auto output = torch::empty({2 * batch_size}, logits.options());
    
    float* loss_ptr = output.data_ptr<float>();
    float* d_logits_ptr = loss_ptr + batch_size;

    const int block_size = 256;
    int num_blocks = (batch_size + block_size - 1) / block_size;
    
    focal_loss_forward_backward_kernel<<<num_blocks, block_size>>>(
        logits.data_ptr<float>(),
        targets.data_ptr<float>(),
        loss_ptr,
        d_logits_ptr,
        batch_size,
        alpha,
        gamma
    );
    
    return output; // 返回包含loss和grads的组合张量
}
"""

# 编译CUDA代码
build_dir = './cuda_build_fused'
os.makedirs(build_dir, exist_ok=True)

focal_loss_module = load_inline(
    name="focal_loss_fused",
    cpp_sources=focal_loss_cpp_source,
    cuda_sources=focal_loss_source,
    functions=["focal_loss_forward_backward_cuda"],
    verbose=True,
    build_directory=build_dir,
    extra_cuda_cflags=["-O3"] 
)

class FocalLossFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, logits, targets, alpha, gamma):
        # 调用融合内核
        output = focal_loss_module.focal_loss_forward_backward_cuda(logits, targets, alpha, gamma)
        
        batch_size = logits.numel()
        loss = output.narrow(0, 0, batch_size)
        d_logits = output.narrow(0, batch_size, batch_size)

        # 为反向传播保存梯度和超参数
        ctx.save_for_backward(d_logits)
        ctx.alpha = alpha
        ctx.gamma = gamma
        return loss

    @staticmethod
    def backward(ctx, grad_output):
        # 从上下文中恢复梯度
        d_logits, = ctx.saved_tensors
        
        # grad_output是上游传来的梯度，需要相乘
        return d_logits * grad_output, None, None, None

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_fn = FocalLossFunction.apply

    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)
        
        # 调用自定义的autograd函数
        loss = self.focal_loss_fn(logits, targets, self.alpha, self.gamma)
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:
            return loss
