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

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

__global__ void focal_loss_labelsmoothing_fused_kernel(
    const float* logits, 
    const float* targets, 
    float* loss, 
    int batch_size,
    float alpha, 
    float gamma,
    float smoothing
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < batch_size) {
        float logit = logits[idx];
        float target = targets[idx];
        
        // Apply label smoothing: smoothed_target = (1-smoothing)*target + smoothing/2
        float smoothed_target = (1.0f - smoothing) * target + smoothing * 0.5f;
        
        // Optimized sigmoid computation with numerical stability
        float prob;
        if (logit > 0.0f) {
            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-8f), 1e-8f);
        
        // Compute pt based on original targets (not smoothed)
        float pt = (target > 0.5f) ? prob : (1.0f - prob);
        
        // Compute focal weight with optimized powf call
        float focal_weight = alpha * powf(1.0f - pt, gamma);
        
        // Optimized binary cross entropy with smoothed targets
        // Direct computation using logits and smoothed targets
        float bce;
        if (logit >= 0.0f) {
            // For positive logits: use log(1 + exp(-logit)) - smoothed_target * logit
            float exp_neg_logit = expf(-logit);
            bce = log1pf(exp_neg_logit) + (1.0f - smoothed_target) * logit;
        } else {
            // For negative logits: -logit + log(1 + exp(logit)) + (1 - smoothed_target) * logit
            float exp_logit = expf(logit);
            bce = -logit + log1pf(exp_logit) + (1.0f - smoothed_target) * logit;
        }
        
        // Apply focal weight
        loss[idx] = focal_weight * bce;
    }
}

torch::Tensor focal_loss_labelsmoothing_fused_cuda(
    torch::Tensor logits, 
    torch::Tensor targets, 
    float alpha, 
    float gamma,
    float smoothing
) {
    // Ensure consistent data types
    logits = logits.to(torch::kFloat32);
    targets = targets.to(torch::kFloat32);
    
    auto batch_size = logits.numel();
    auto loss = torch::empty_like(logits);
    
    // Optimized block configuration
    const int block_size = 256;
    int num_blocks = (batch_size + block_size - 1) / block_size;
    
    // Launch optimized kernel
    focal_loss_labelsmoothing_fused_kernel<<<num_blocks, block_size>>>(
        logits.data_ptr<float>(),
        targets.data_ptr<float>(),
        loss.data_ptr<float>(),
        batch_size,
        alpha,
        gamma,
        smoothing
    );
    
    // Check for CUDA errors
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        AT_ERROR("CUDA kernel failed: ", cudaGetErrorString(err));
    }
    
    return loss;
}
"""

focal_loss_labelsmoothing_cpp_source = """
torch::Tensor focal_loss_labelsmoothing_fused_cuda(
    torch::Tensor logits, 
    torch::Tensor targets, 
    float alpha, 
    float gamma,
    float smoothing
);
"""

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

# Compile the optimized inline CUDA code
focal_loss_labelsmoothing_module = load_inline(
    name="focal_loss_labelsmoothing_fused",
    cpp_sources=focal_loss_labelsmoothing_cpp_source,
    cuda_sources=focal_loss_labelsmoothing_source,
    functions=["focal_loss_labelsmoothing_fused_cuda"],
    verbose=True,
    build_directory=build_dir,
    extra_cuda_cflags=["-O3", "--use_fast_math"]
)

class ModelNew(torch.nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0, smoothing=0.1, reduction='mean'):
        super(ModelNew, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.smoothing = smoothing
        self.reduction = reduction
        self.focal_loss_labelsmoothing = focal_loss_labelsmoothing_module

    def forward(self, logits, targets):
        # Ensure consistent input types
        logits = logits.to(torch.float32)
        targets = targets.to(torch.float32)
        
        # Handle shape matching
        if logits.dim() == 2 and logits.size(1) == 1:
            logits = logits.squeeze(1)
        
        # Compute focal loss with label smoothing using optimized fused CUDA kernel
        loss = self.focal_loss_labelsmoothing.focal_loss_labelsmoothing_fused_cuda(
            logits, targets, self.alpha, self.gamma, self.smoothing
        )
        
        # Apply reduction
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:
            return loss
