import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# C++ 源代码 wrapper
cpp_source = """
#include <torch/extension.h>
#include <string>

torch::Tensor asl_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& targets,
    float gamma_pos,
    float gamma_neg,
    float margin,
    float eps,
    std::string reduction);
"""

# CUDA 源代码
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

#define BLOCK_SIZE 256

struct __align__(16) Float4 {
    float x, y, z, w;
};

// Device Helper
__device__ __forceinline__ float compute_asl_point(
    float logit, float target, 
    float g_pos, float g_neg, float margin, float eps) 
{
    float p = 1.0f / (1.0f + __expf(-logit));
    
    // Branchless selection usually better, but here branches are aligned (target is integer 0/1)
    if (target > 0.5f) {
        // Positive case: y = 1
        // Loss = -(1-p)^g_pos * log(p)
        float coeff = __powf(1.0f - p, g_pos);
        return -coeff * __logf(p + eps);
    } else {
        // Negative case: y = 0
        // p_m = max(p - margin, 0)
        float p_m = p - margin;
        if (p_m <= 0.0f) {
            return 0.0f; // Hard thresholding optimization
        }
        // Loss = -(p_m)^g_neg * log(1 - p_m)
        float coeff = __powf(p_m, g_neg);
        return -coeff * __logf(1.0f - p_m + eps);
    }
}

__global__ void asl_loss_kernel(
    float* __restrict__ output,
    const float* __restrict__ logits,
    const float* __restrict__ targets,
    const int n_elements,
    const float g_pos,
    const float g_neg,
    const float margin,
    const float eps) 
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int vec_n = n_elements / 4;
    int stride = blockDim.x * gridDim.x;

    // 1. Vectorized Loop
    for (int i = idx; i < vec_n; i += stride) {
        Float4 l_vec = reinterpret_cast<const Float4*>(logits)[i];
        Float4 t_vec = reinterpret_cast<const Float4*>(targets)[i];
        Float4 out_vec;

        out_vec.x = compute_asl_point(l_vec.x, t_vec.x, g_pos, g_neg, margin, eps);
        out_vec.y = compute_asl_point(l_vec.y, t_vec.y, g_pos, g_neg, margin, eps);
        out_vec.z = compute_asl_point(l_vec.z, t_vec.z, g_pos, g_neg, margin, eps);
        out_vec.w = compute_asl_point(l_vec.w, t_vec.w, g_pos, g_neg, margin, eps);

        reinterpret_cast<Float4*>(output)[i] = out_vec;
    }

    // 2. Scalar Tail
    int tail_start = vec_n * 4;
    int global_tid = idx * 4; // Approx mapping
    // Simple absolute mapping for tail
    int scalar_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (scalar_idx < (n_elements - tail_start)) {
        int real_idx = tail_start + scalar_idx;
        output[real_idx] = compute_asl_point(
            logits[real_idx], targets[real_idx], g_pos, g_neg, margin, eps);
    }
}

torch::Tensor asl_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& targets,
    float g_pos,
    float g_neg,
    float margin,
    float eps,
    std::string reduction) 
{
    TORCH_CHECK(logits.is_cuda() && targets.is_cuda(), "Inputs must be CUDA");
    TORCH_CHECK(logits.is_contiguous() && targets.is_contiguous(), "Inputs must be contiguous");
    TORCH_CHECK(logits.numel() == targets.numel(), "Shapes must match");

    const int n = logits.numel();
    auto output = torch::empty_like(logits);

    const int vec_n = n / 4;
    const int grid_size = (vec_n + BLOCK_SIZE - 1) / BLOCK_SIZE;
    int final_grid = (grid_size < 1) ? 1 : grid_size;
    if (final_grid > 65535) final_grid = 65535;

    asl_loss_kernel<<<final_grid, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        logits.data_ptr<float>(),
        targets.data_ptr<float>(),
        n,
        g_pos,
        g_neg,
        margin,
        eps
    );

    if (reduction == "mean") {
        return output.mean();
    } else if (reduction == "sum") {
        return output.sum();
    }

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, gamma_neg=4.0, gamma_pos=1.0, clip=0.05, reduction='none'):
        super(ModelNew, self).__init__()
        self.gamma_neg = gamma_neg
        self.gamma_pos = gamma_pos
        self.clip = clip
        self.reduction = reduction
        self.eps = 1e-8
        
        self.op = load_inline(
            name='asl_loss_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['asl_loss_cuda_forward'],
            verbose=False,
            extra_cuda_cflags=['-O3']
        )

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        return self.op.asl_loss_cuda_forward(
            logits.contiguous(), 
            targets.contiguous(), 
            self.gamma_pos,
            self.gamma_neg,
            self.clip,
            self.eps,
            self.reduction
        )