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

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

torch::Tensor wing_loss_cuda_forward(
    const torch::Tensor& pred, 
    const torch::Tensor& target,
    float w, 
    float epsilon,
    float C,
    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__ __forceinline__ float compute_point(float p, float t, float w, float eps, float C) {
    float dist = fabsf(t - p);
    if (dist < w) {
        // w * ln(1 + |x|/eps)
        return w * logf(1.0f + dist / eps);
    } else {
        // |x| - C
        return dist - C;
    }
}

__global__ void wing_loss_kernel(
    float* __restrict__ output,
    const float* __restrict__ pred,
    const float* __restrict__ target,
    const int n_elements,
    const float w,
    const float eps,
    const float C) 
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;

    const int vec_n = n_elements / 4;
    
    // Grid-Stride Loop (以 float4 为步长)
    int i = idx;
    const int stride = blockDim.x * gridDim.x;

    for (; i < vec_n; i += stride) {
        Float4 p_vec = reinterpret_cast<const Float4*>(pred)[i];
        Float4 t_vec = reinterpret_cast<const Float4*>(target)[i];
        Float4 out_vec;

        out_vec.x = compute_point(p_vec.x, t_vec.x, w, eps, C);
        out_vec.y = compute_point(p_vec.y, t_vec.y, w, eps, C);
        out_vec.z = compute_point(p_vec.z, t_vec.z, w, eps, C);
        out_vec.w = compute_point(p_vec.w, t_vec.w, w, eps, C);

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

    // 2. Scalar Part (Tail)
    int scalar_i = i * 4; 
    
    int tail_start = vec_n * 4;
    int global_tid = idx; 
    
    if (global_tid < (n_elements - tail_start)) {
        int real_idx = tail_start + global_tid;
        output[real_idx] = compute_point(pred[real_idx], target[real_idx], w, eps, C);
    }
}

torch::Tensor wing_loss_cuda_forward(
    const torch::Tensor& pred, 
    const torch::Tensor& target,
    float w, 
    float epsilon,
    float C,
    std::string reduction) 
{
    TORCH_CHECK(pred.is_cuda() && target.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(pred.is_contiguous() && target.is_contiguous(), "Inputs must be contiguous");
    TORCH_CHECK(pred.numel() == target.numel(), "Shapes must match");

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

    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;
    
    // 限制最大 Grid 以防过大
    if (final_grid > 65535) final_grid = 65535; 

    wing_loss_kernel<<<final_grid, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        pred.data_ptr<float>(),
        target.data_ptr<float>(),
        n,
        w,
        epsilon,
        C
    );

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

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, w=10.0, epsilon=2.0, reduction='none'):
        super(ModelNew, self).__init__()
        self.w = w
        self.epsilon = epsilon
        self.reduction = reduction
        # Pre-compute Constant C
        self.C = self.w - self.w * math.log(1 + self.w / self.epsilon)
        
        self.op = load_inline(
            name='wing_loss_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['wing_loss_cuda_forward'],
            verbose=False,
            extra_cuda_cflags=['-O3']
        )

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        return self.op.wing_loss_cuda_forward(
            pred.contiguous(), 
            target.contiguous(), 
            self.w, 
            self.epsilon,
            self.C,
            self.reduction
        )