import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
import math
cpp_source = """
#include <torch/extension.h>
#include <string>

torch::Tensor huber_loss_cuda_forward(
    const torch::Tensor& input, 
    const torch::Tensor& target,
    double delta,
    const std::string& reduction
);
"""

# CUDA 源代码，包含核函数和其调用封装
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <limits>

#define BLOCK_SIZE 256

// ----------------------------------------------------------------------------
// Element-wise Kernel for reduction='none'
// ----------------------------------------------------------------------------
template <typename T>
__global__ void huber_loss_elementwise_kernel(
    T* output,
    const T* input,
    const T* target,
    T delta,
    int64_t n_elements) 
{
    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n_elements) return;

    T diff = input[idx] - target[idx];
    T abs_diff = fabsf(diff);
    if (abs_diff < delta) {
        output[idx] = 0.5f * diff * diff;
    } else {
        output[idx] = delta * (abs_diff - 0.5f * delta);
    }
}


// ----------------------------------------------------------------------------
// Reduction Kernel: Stage 1 (calculate loss and reduce within blocks)
// ----------------------------------------------------------------------------
template <typename T>
__global__ void huber_loss_reduce_kernel_stage1(
    T* block_results,
    const T* input,
    const T* target,
    T delta,
    int64_t n_elements)
{
    __shared__ T sdata[BLOCK_SIZE];
    
    int64_t tid = threadIdx.x;
    int64_t i = blockIdx.x * (blockDim.x * 2) + tid;
    int64_t gridSize = blockDim.x * 2 * gridDim.x;

    T my_sum = 0.0f;

    while (i < n_elements) {
        T diff1 = input[i] - target[i];
        T abs_diff1 = fabsf(diff1);
        my_sum += (abs_diff1 < delta) ? (0.5f * diff1 * diff1) : (delta * (abs_diff1 - 0.5f * delta));

        if (i + blockDim.x < n_elements) {
            T diff2 = input[i + blockDim.x] - target[i + blockDim.x];
            T abs_diff2 = fabsf(diff2);
            my_sum += (abs_diff2 < delta) ? (0.5f * diff2 * diff2) : (delta * (abs_diff2 - 0.5f * delta));
        }
        i += gridSize;
    }
    sdata[tid] = my_sum;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    // 第一个线程将该块的结果写入全局内存
    if (tid == 0) {
        block_results[blockIdx.x] = sdata[0];
    }
}


torch::Tensor huber_loss_cuda_forward(
    const torch::Tensor& input, 
    const torch::Tensor& target,
    double delta_d,
    const std::string& reduction)
{
    TORCH_CHECK(input.is_cuda(), "Input tensor must be a CUDA tensor");
    TORCH_CHECK(target.is_cuda(), "Target tensor must be a CUDA tensor");
    TORCH_CHECK(input.sizes() == target.sizes(), "Input and target shapes must match");
    TORCH_CHECK(input.is_contiguous(), "Input tensor must be contiguous");
    TORCH_CHECK(target.is_contiguous(), "Target tensor must be contiguous");

    const int64_t n_elements = input.numel();
    auto scalar_type = input.scalar_type();

    if (reduction == "none") {
        auto output = torch::empty_like(input);
        const int num_blocks = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
        AT_DISPATCH_FLOATING_TYPES(scalar_type, "huber_loss_elementwise", ([&] {
            huber_loss_elementwise_kernel<scalar_t><<<num_blocks, BLOCK_SIZE>>>(
                output.data_ptr<scalar_t>(),
                input.data_ptr<scalar_t>(),
                target.data_ptr<scalar_t>(),
                static_cast<scalar_t>(delta_d),
                n_elements);
        }));
        return output;
    } 
    else // 'sum' or 'mean'
    {
        int max_grid_size = 1024; // A reasonable limit for partial results
        int num_blocks = std::min(max_grid_size, (int)((n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE));
        
        auto options = torch::TensorOptions().device(input.device()).dtype(input.dtype());
        auto block_results = torch::empty({num_blocks}, options);
        auto output = torch::empty({}, options);

        AT_DISPATCH_FLOATING_TYPES(scalar_type, "huber_loss_reduce_stage1", ([&] {
            huber_loss_reduce_kernel_stage1<scalar_t><<<num_blocks, BLOCK_SIZE>>>(
                block_results.data_ptr<scalar_t>(),
                input.data_ptr<scalar_t>(),
                target.data_ptr<scalar_t>(),
                static_cast<scalar_t>(delta_d),
                n_elements);
        }));

        // Stage 2: final reduction of block_results.
        // For simplicity and robustness, we can sum the small block_results tensor on the CPU
        // or launch another kernel. A simple sum() is often fast enough here.
        // A full GPU stage2 implementation would be similar to stage1 but on `block_results`.
        torch::Tensor total_sum_tensor = block_results.sum();
        
        if (reduction == "mean") {
            return total_sum_tensor / n_elements;
        }
        return total_sum_tensor;
    }
}

"""

class ModelNew(nn.Module):
    """
    使用自定义 CUDA 内核进行优化的 HuberLoss 模型。
    """
    def __init__(self, delta=1.0, reduction='mean'):
        super(ModelNew, self).__init__()
        self.delta = delta
        self.reduction = reduction
        
        self.huber_loss_op = load_inline(
            name='huber_loss_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['huber_loss_cuda_forward'],
            verbose=False,
            extra_cuda_cflags=["-O3"]
        )

    def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
        return self.huber_loss_op.huber_loss_cuda_forward(
            input_tensor, 
            target_tensor,
            self.delta,
            self.reduction
        )