import torch
from torch.utils.cpp_extension import load_inline

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

// 基础版本 - 分步计算
__global__ void cosinedistance_kernel_basic(
    const float* __restrict__ x,
    const float* __restrict__ y,
    float* __restrict__ distances,
    int batch_size,
    int feature_dim,
    float eps
) {
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    int base = sample_idx * feature_dim;
    
    // 计算点积
    float dot_product = 0.0f;
    float norm_x_sq = 0.0f;
    float norm_y_sq = 0.0f;
    
    // 4元素展开计算
    int dim = 0;
    for (; dim <= feature_dim - 4; dim += 4) {
        float x0 = x[base + dim];
        float x1 = x[base + dim + 1];
        float x2 = x[base + dim + 2];
        float x3 = x[base + dim + 3];
        
        float y0 = y[base + dim];
        float y1 = y[base + dim + 1];
        float y2 = y[base + dim + 2];
        float y3 = y[base + dim + 3];
        
        // 点积
        dot_product += x0*y0 + x1*y1 + x2*y2 + x3*y3;
        
        // 范数平方
        norm_x_sq += x0*x0 + x1*x1 + x2*x2 + x3*x3;
        norm_y_sq += y0*y0 + y1*y1 + y2*y2 + y3*y3;
    }
    
    // 处理剩余元素
    for (; dim < feature_dim; dim++) {
        float x_val = x[base + dim];
        float y_val = y[base + dim];
        
        dot_product += x_val * y_val;
        norm_x_sq += x_val * x_val;
        norm_y_sq += y_val * y_val;
    }
    
    // 计算余弦距离
    float norm_x = sqrtf(norm_x_sq);
    float norm_y = sqrtf(norm_y_sq);
    float norm_product = norm_x * norm_y;
    
    // 避免除零
    if (norm_product < eps) {
        norm_product = eps;
    }
    
    float cosine_sim = dot_product / norm_product;
    distances[sample_idx] = 1.0f - cosine_sim;
}

// 向量化版本 - float4优化
__global__ void cosinedistance_kernel_vectorized(
    const float* __restrict__ x,
    const float* __restrict__ y,
    float* __restrict__ distances,
    int batch_size,
    int feature_dim,
    float eps
) {
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    // 使用float4向量化
    const float4* x_vec = reinterpret_cast<const float4*>(x);
    const float4* y_vec = reinterpret_cast<const float4*>(y);
    
    int feature_dim_vec = feature_dim / 4;
    int base_vec = sample_idx * feature_dim_vec;
    
    float dot_product = 0.0f;
    float norm_x_sq = 0.0f;
    float norm_y_sq = 0.0f;
    
    // 向量化计算
    for (int dim_vec = 0; dim_vec < feature_dim_vec; dim_vec++) {
        float4 x_val = x_vec[base_vec + dim_vec];
        float4 y_val = y_vec[base_vec + dim_vec];
        
        // 点积
        dot_product += x_val.x*y_val.x + x_val.y*y_val.y + x_val.z*y_val.z + x_val.w*y_val.w;
        
        // 范数平方
        norm_x_sq += x_val.x*x_val.x + x_val.y*x_val.y + x_val.z*x_val.z + x_val.w*x_val.w;
        norm_y_sq += y_val.x*y_val.x + y_val.y*y_val.y + y_val.z*y_val.z + y_val.w*y_val.w;
    }
    
    // 处理剩余元素
    for (int dim = feature_dim_vec * 4; dim < feature_dim; dim++) {
        float x_val = x[sample_idx * feature_dim + dim];
        float y_val = y[sample_idx * feature_dim + dim];
        
        dot_product += x_val * y_val;
        norm_x_sq += x_val * x_val;
        norm_y_sq += y_val * y_val;
    }
    
    // 计算余弦距离
    float norm_x = sqrtf(norm_x_sq);
    float norm_y = sqrtf(norm_y_sq);
    float norm_product = norm_x * norm_y;
    
    // 避免除零
    if (norm_product < eps) {
        norm_product = eps;
    }
    
    float cosine_sim = dot_product / norm_product;
    distances[sample_idx] = 1.0f - cosine_sim;
}

// 融合版本 - 预计算范数
__global__ void cosinedistance_kernel_fused(
    const float* __restrict__ x,
    const float* __restrict__ y,
    const float* __restrict__ x_norms,
    const float* __restrict__ y_norms,
    float* __restrict__ distances,
    int batch_size,
    int feature_dim,
    float eps
) {
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    // 使用float4向量化
    const float4* x_vec = reinterpret_cast<const float4*>(x);
    const float4* y_vec = reinterpret_cast<const float4*>(y);
    
    int feature_dim_vec = feature_dim / 4;
    int base_vec = sample_idx * feature_dim_vec;
    
    float dot_product = 0.0f;
    
    // 只计算点积，范数已经预计算
    for (int dim_vec = 0; dim_vec < feature_dim_vec; dim_vec++) {
        float4 x_val = x_vec[base_vec + dim_vec];
        float4 y_val = y_vec[base_vec + dim_vec];
        
        dot_product += x_val.x*y_val.x + x_val.y*y_val.y + x_val.z*y_val.z + x_val.w*y_val.w;
    }
    
    // 处理剩余元素
    for (int dim = feature_dim_vec * 4; dim < feature_dim; dim++) {
        dot_product += x[sample_idx * feature_dim + dim] * y[sample_idx * feature_dim + dim];
    }
    
    // 使用预计算的范数
    float norm_product = x_norms[sample_idx] * y_norms[sample_idx];
    
    // 避免除零
    if (norm_product < eps) {
        norm_product = eps;
    }
    
    float cosine_sim = dot_product / norm_product;
    distances[sample_idx] = 1.0f - cosine_sim;
}

// 预计算范数kernel
__global__ void compute_norms_kernel(
    const float* __restrict__ data,
    float* __restrict__ norms,
    int batch_size,
    int feature_dim
) {
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    int base = sample_idx * feature_dim;
    float norm_sq = 0.0f;
    
    // 4元素展开计算范数平方
    int dim = 0;
    for (; dim <= feature_dim - 4; dim += 4) {
        float val0 = data[base + dim];
        float val1 = data[base + dim + 1];
        float val2 = data[base + dim + 2];
        float val3 = data[base + dim + 3];
        
        norm_sq += val0*val0 + val1*val1 + val2*val2 + val3*val3;
    }
    
    // 处理剩余元素
    for (; dim < feature_dim; dim++) {
        float val = data[base + dim];
        norm_sq += val * val;
    }
    
    norms[sample_idx] = sqrtf(norm_sq);
}

torch::Tensor cosinedistance_cuda(
    torch::Tensor x,
    torch::Tensor y,
    float eps,
    std::string mode = "vectorized"
) {
    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());
    
    if (mode == "fused") {
        // 预计算范数
        auto x_norms = torch::zeros({batch_size}, x.options());
        auto y_norms = torch::zeros({batch_size}, y.options());
        
        compute_norms_kernel<<<batch_size, 1>>>(
            x_contig.data_ptr<float>(),
            x_norms.data_ptr<float>(),
            batch_size,
            feature_dim
        );
        
        compute_norms_kernel<<<batch_size, 1>>>(
            y_contig.data_ptr<float>(),
            y_norms.data_ptr<float>(),
            batch_size,
            feature_dim
        );
        
        // 融合版本
        cosinedistance_kernel_fused<<<batch_size, 1>>>(
            x_contig.data_ptr<float>(),
            y_contig.data_ptr<float>(),
            x_norms.data_ptr<float>(),
            y_norms.data_ptr<float>(),
            distances.data_ptr<float>(),
            batch_size,
            feature_dim,
            eps
        );
    } else if (mode == "vectorized") {
        // 向量化版本（默认）
        cosinedistance_kernel_vectorized<<<batch_size, 1>>>(
            x_contig.data_ptr<float>(),
            y_contig.data_ptr<float>(),
            distances.data_ptr<float>(),
            batch_size,
            feature_dim,
            eps
        );
    } else {
        // 基础版本
        cosinedistance_kernel_basic<<<batch_size, 1>>>(
            x_contig.data_ptr<float>(),
            y_contig.data_ptr<float>(),
            distances.data_ptr<float>(),
            batch_size,
            feature_dim,
            eps
        );
    }
    
    return distances;
}
"""

cosinedistance_cpp_source = """
torch::Tensor cosinedistance_cuda(torch::Tensor x, torch::Tensor y, float eps, std::string mode);
"""

# 编译CUDA代码
cosinedistance = load_inline(
    name="cosinedistance",
    cpp_sources=cosinedistance_cpp_source,
    cuda_sources=cosinedistance_source,
    functions=["cosinedistance_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, eps=1e-8, mode="vectorized"):
        super(ModelNew, self).__init__()
        self.eps = eps
        self.mode = mode
        self.cosinedistance = cosinedistance

    def forward(self, x, y):
        return self.cosinedistance.cosinedistance_cuda(x, y, self.eps, self.mode)
