import torch
from torch.utils.cpp_extension import load_inline

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


__global__ void braycurtis_kernel_vectorized(
    const float* __restrict__ x,
    const float* __restrict__ y,
    float* __restrict__ distances,
    int batch_size,
    int feature_dim
) {
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    int tid = threadIdx.x;
    int base = sample_idx * feature_dim;
    
    // 每个线程独立计算自己的部分，然后归约
    float local_num = 0.0f;
    float local_den = 0.0f;
    
    // 每个线程处理4个元素（float4向量化）
    int stride = blockDim.x * 4;
    for (int dim = tid * 4; dim < feature_dim; dim += stride) {
        // 确保不越界
        if (dim + 3 < feature_dim) {
            float4 x_val = *reinterpret_cast<const float4*>(&x[base + dim]);
            float4 y_val = *reinterpret_cast<const float4*>(&y[base + dim]);
            
            // 处理4个分量 - Bray-Curtis计算
            float abs_diff1 = fabsf(x_val.x - y_val.x);
            float abs_sum1 = fabsf(x_val.x) + fabsf(y_val.x);
            local_num += abs_diff1;
            local_den += abs_sum1;
            
            float abs_diff2 = fabsf(x_val.y - y_val.y);
            float abs_sum2 = fabsf(x_val.y) + fabsf(y_val.y);
            local_num += abs_diff2;
            local_den += abs_sum2;
            
            float abs_diff3 = fabsf(x_val.z - y_val.z);
            float abs_sum3 = fabsf(x_val.z) + fabsf(y_val.z);
            local_num += abs_diff3;
            local_den += abs_sum3;
            
            float abs_diff4 = fabsf(x_val.w - y_val.w);
            float abs_sum4 = fabsf(x_val.w) + fabsf(y_val.w);
            local_num += abs_diff4;
            local_den += abs_sum4;
        } else {
            // 处理剩余元素
            for (int i = dim; i < feature_dim; i++) {
                float x_val = x[base + i];
                float y_val = y[base + i];
                
                float abs_diff = fabsf(x_val - y_val);
                float abs_sum = fabsf(x_val) + fabsf(y_val);
                
                local_num += abs_diff;
                local_den += abs_sum;
            }
        }
    }
    
    // 使用共享内存进行归约
    extern __shared__ float shared_data[];
    float* shared_num = shared_data;
    float* shared_den = shared_data + blockDim.x;
    
    shared_num[tid] = local_num;
    shared_den[tid] = local_den;
    
    __syncthreads();
    
    // 归约求和 - 分子和分母分别归约
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            shared_num[tid] += shared_num[tid + stride];
            shared_den[tid] += shared_den[tid + stride];
        }
        __syncthreads();
    }
    
    // 第一个线程写入结果
    if (tid == 0) {
        if (shared_den[0] > 0.0f) {
            distances[sample_idx] = shared_num[0] / shared_den[0];
        } else {
            distances[sample_idx] = 0.0f;
        }
    }
}

torch::Tensor braycurtis_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");
    
    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 = 64;
    size_t shared_mem = block_size * 2 * sizeof(float);
    
    braycurtis_kernel_vectorized<<<batch_size, block_size, shared_mem>>>(
        x_contig.data_ptr<float>(),
        y_contig.data_ptr<float>(),
        distances.data_ptr<float>(),
        batch_size,
        feature_dim
    );
    
    return distances;
}
"""

braycurtis_cpp_source = """
torch::Tensor braycurtis_cuda(torch::Tensor x, torch::Tensor y);
"""

# Compile the inline CUDA code
braycurtis = load_inline(
    name="braycurtis",
    cpp_sources=braycurtis_cpp_source,
    cuda_sources=braycurtis_source,
    functions=["braycurtis_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):
        super(ModelNew, self).__init__()
        self.braycurtis = braycurtis

    def forward(self, x, y):
        return self.braycurtis.braycurtis_cuda(x, y)
