import torch
from torch.utils.cpp_extension import load_inline

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

// 多线程版本 - 通用p值
__global__ void minkowski_kernel_general(
    const float* __restrict__ x,
    const float* __restrict__ y,
    float* __restrict__ distances,
    int batch_size,
    int feature_dim,
    float p,
    float inv_p
) {
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    int tid = threadIdx.x;
    int base = sample_idx * feature_dim;
    
    // 使用共享内存进行归约
    extern __shared__ float shared_sum[];
    shared_sum[tid] = 0.0f;
    
    // 每个线程处理多个元素
    int stride = blockDim.x;
    for (int dim = tid; dim < feature_dim; dim += stride) {
        float diff = x[base + dim] - y[base + dim];
        float abs_diff = fabsf(diff);
        shared_sum[tid] += powf(abs_diff, p);
    }
    
    __syncthreads();
    
    // 块内归约求和
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            shared_sum[tid] += shared_sum[tid + stride];
        }
        __syncthreads();
    }
    
    // 第一个线程计算最终结果
    if (tid == 0) {
        distances[sample_idx] = powf(shared_sum[0], inv_p);
    }
}

// 特殊优化版本 - p=2 (欧几里得距离)
__global__ void minkowski_kernel_p2(
    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 base = sample_idx * feature_dim;
    float sum_sq = 0.0f;
    
    // 计算Σ(x_i - y_i)²
    for (int dim = 0; dim < feature_dim; dim++) {
        float diff = x[base + dim] - y[base + dim];
        sum_sq += diff * diff;
    }
    
    // 计算√Σ(x_i - y_i)²
    distances[sample_idx] = sqrtf(sum_sq);
}

// 特殊优化版本 - p=1 (曼哈顿距离)
__global__ void minkowski_kernel_p1(
    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 base = sample_idx * feature_dim;
    float sum_abs = 0.0f;
    
    // 计算Σ|x_i - y_i|
    for (int dim = 0; dim < feature_dim; dim++) {
        float diff = x[base + dim] - y[base + dim];
        sum_abs += fabsf(diff);
    }
    
    distances[sample_idx] = sum_abs;
}

torch::Tensor minkowski_cuda(torch::Tensor x, torch::Tensor y, float p) {
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "X must be float32");
    TORCH_CHECK(y.scalar_type() == torch::kFloat32, "Y must be float32");
    TORCH_CHECK(p > 0, "p must be positive");
    
    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());
    
    if (p == 1.0f) {
        // 特殊优化：曼哈顿距离
        minkowski_kernel_p1<<<batch_size, 1>>>(
            x_contig.data_ptr<float>(),
            y_contig.data_ptr<float>(),
            distances.data_ptr<float>(),
            batch_size,
            feature_dim
        );
    } else if (p == 2.0f) {
        // 特殊优化：欧几里得距离
        minkowski_kernel_p2<<<batch_size, 1>>>(
            x_contig.data_ptr<float>(),
            y_contig.data_ptr<float>(),
            distances.data_ptr<float>(),
            batch_size,
            feature_dim
        );
    } else {
        // 通用情况 - 多线程版本
        float inv_p = 1.0f / p;
        const int block_size = 256;
        size_t shared_mem = block_size * sizeof(float);
        
        minkowski_kernel_general<<<batch_size, block_size, shared_mem>>>(
            x_contig.data_ptr<float>(),
            y_contig.data_ptr<float>(),
            distances.data_ptr<float>(),
            batch_size,
            feature_dim,
            p,
            inv_p
        );
    }
    
    return distances;
}
"""

minkowski_cpp_source = """
torch::Tensor minkowski_cuda(torch::Tensor x, torch::Tensor y, float p);
"""

# Compile the inline CUDA code
minkowski = load_inline(
    name="minkowski",
    cpp_sources=minkowski_cpp_source,
    cuda_sources=minkowski_source,
    functions=["minkowski_cuda"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, p=2):
        super(ModelNew, self).__init__()
        self.p = p
        self.minkowski = minkowski  # The module containing the kernel

    def forward(self, x, y):
        return self.minkowski.minkowski_cuda(x, y, self.p)
