import torch
from torch.utils.cpp_extension import load_inline

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

// 高效归约版本 - 最优性能和精度
__global__ void huberloss_kernel_optimal(
    const float* __restrict__ input,
    const float* __restrict__ target,
    float* __restrict__ partial_sums,
    int total_elements,
    float delta
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 局部累加器
    float local_loss = 0.0f;
    
    // 每个线程处理多个元素
    int stride = blockDim.x * gridDim.x;
    for (int element_idx = idx; element_idx < total_elements; element_idx += stride) {
        // 计算误差
        float error = input[element_idx] - target[element_idx];
        float abs_error = fabsf(error);
        
        // Huber Loss公式
        float loss;
        if (abs_error <= delta) {
            loss = 0.5f * error * error;
        } else {
            loss = delta * (abs_error - 0.5f * delta);
        }
        
        local_loss += loss;
    }
    
    // 使用共享内存进行块内归约
    extern __shared__ float shared_mem[];
    shared_mem[threadIdx.x] = local_loss;
    
    __syncthreads();
    
    // 块内归约
    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
        if (threadIdx.x < stride) {
            shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
        }
        __syncthreads();
    }
    
    // 每个block写入部分和到全局内存
    if (threadIdx.x == 0) {
        partial_sums[blockIdx.x] = shared_mem[0];
    }
}

torch::Tensor huberloss_cuda(
    torch::Tensor input, 
    torch::Tensor target, 
    float delta, 
    std::string mode = "optimal"
) {
    TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be float32");
    TORCH_CHECK(target.scalar_type() == torch::kFloat32, "Target must be float32");
    TORCH_CHECK(input.sizes() == target.sizes(), "Input and target must have same shape");
    
    auto input_contig = input.contiguous();
    auto target_contig = target.contiguous();
    
    int total_elements = input_contig.numel();
    
    // 只使用最优的efficient模式
    const int block_size = 256;
    
    // 针对不同数据量的优化block分配
    int num_blocks;
    if (total_elements <= 2048) {
        num_blocks = 4;
    } else if (total_elements <= 8192) {
        num_blocks = 8;
    } else if (total_elements <= 32768) {
        num_blocks = 16;
    } else if (total_elements <= 131072) {
        num_blocks = 32;
    } else {
        num_blocks = 64;
    }
    
    // 创建部分和数组
    auto partial_sums = torch::zeros({num_blocks}, input.options());
    
    size_t shared_mem = block_size * sizeof(float);
    huberloss_kernel_optimal<<<num_blocks, block_size, shared_mem>>>(
        input_contig.data_ptr<float>(),
        target_contig.data_ptr<float>(),
        partial_sums.data_ptr<float>(),
        total_elements,
        delta
    );
    
    // 在GPU上完成最终归约
    auto total_loss = torch::sum(partial_sums);
    return total_loss;
}
"""

huberloss_cpp_source = """
torch::Tensor huberloss_cuda(torch::Tensor input, torch::Tensor target, float delta, std::string mode);
"""

# 编译CUDA代码 - 平衡精度和性能
huberloss = load_inline(
    name="huberloss",
    cpp_sources=huberloss_cpp_source,
    cuda_sources=huberloss_source,
    functions=["huberloss_cuda"],
    extra_cuda_cflags=[
        "-O3", 
        "--use_fast_math",
        "-gencode=arch=compute_80,code=sm_80"
    ],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, delta=1.0, mode="optimal"):
        super(ModelNew, self).__init__()
        self.delta = delta
        self.mode = mode
        self.huberloss = huberloss

    def forward(self, input, target):
        return self.huberloss.huberloss_cuda(input, target, self.delta, self.mode)
