import torch
from torch.utils.cpp_extension import load_inline

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

#define BLOCK_SIZE 256

// 每个线程块处理一个样本，块内线程并行处理特征维度
__global__ void canberra_distance_kernel(
    const float* __restrict__ x,
    const float* __restrict__ y,
    float* __restrict__ distances,
    int feature_dim
) {
    // 使用动态共享内存进行块内归约
    extern __shared__ float sdata[];
    
    int tid = threadIdx.x;
    int sample_idx = blockIdx.x;
    
    // 每个线程计算自己负责的部分和
    float partial_sum = 0.0f;
    int base_addr = sample_idx * feature_dim;
    
    // 循环处理所有特征，步长为线程数
    for (int i = tid; i < feature_dim; i += blockDim.x) {
        float x_val = x[base_addr + i];
        float y_val = y[base_addr + i];
        
        float abs_diff = fabsf(x_val - y_val);
        float abs_sum = fabsf(x_val) + fabsf(y_val);
        
        if (abs_sum > 0.0f) {
            partial_sum += abs_diff / abs_sum;
        }
    }
    
    sdata[tid] = partial_sum;
    __syncthreads();
    
    // 标准的并行归约求和
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            sdata[tid] += sdata[tid + stride];
        }
        __syncthreads();
    }
    
    // 线程0将最终结果写回
    if (tid == 0) {
        distances[sample_idx] = sdata[0];
    }
}

torch::Tensor canberra_cuda(torch::Tensor x, torch::Tensor y) {
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "X must be float32");
    TORCH_CHECK(y.scalar_type() == torch::kFloat32, "Y must be float32");
    TORCH_CHECK(x.sizes() == y.sizes(), "X and Y must have the same shape");

    auto x_contig = x.contiguous();
    auto y_contig = y.contiguous();
    
    int batch_size = x_contig.size(0);
    int feature_dim = x_contig.size(1);
    
    auto distances = torch::zeros({batch_size}, x.options());
    
    const int block_size = BLOCK_SIZE;
    size_t shared_mem = block_size * sizeof(float);
    
    canberra_distance_kernel<<<batch_size, block_size, shared_mem>>>(
        x_contig.data_ptr<float>(),
        y_contig.data_ptr<float>(),
        distances.data_ptr<float>(),
        feature_dim
    );
    
    return distances;
}
"""

canberra_cpp_source = """
torch::Tensor canberra_cuda(torch::Tensor x, torch::Tensor y);
"""

# 编译CUDA代码
canberra = load_inline(
    name="canberra_fixed",
    cpp_sources=canberra_cpp_source,
    cuda_sources=canberra_source,
    functions=["canberra_cuda"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.canberra = canberra

    def forward(self, x, y):
        return self.canberra.canberra_cuda(x, y)
