import torch
from torch.utils.cpp_extension import load_inline

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

// 基础版本 - 简单比较
__global__ void hamming_kernel_basic(
    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;
    int distance = 0;
    
    // Count different elements
    for (int i = 0; i < feature_dim; i++) {
        if (x[base + i] != y[base + i]) {
            distance++;
        }
    }
    
    distances[sample_idx] = distance;
}

// 向量化版本 - float4优化
__global__ void hamming_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;
    
    // 使用共享内存进行归约
    extern __shared__ float shared_sum[];
    shared_sum[tid] = 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]);
            
            // 比较每个分量
            shared_sum[tid] += (x_val.x != y_val.x) + (x_val.y != y_val.y) + 
                              (x_val.z != y_val.z) + (x_val.w != y_val.w);
        } else {
            // 处理剩余元素
            for (int i = dim; i < feature_dim; i++) {
                if (x[base + i] != y[base + i]) {
                    shared_sum[tid] += 1.0f;
                }
            }
        }
    }
    
    __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] = shared_sum[0];
    }
}

// 融合版本 - 单kernel完成所有计算
__global__ void hamming_kernel_fused(
    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;
    
    // 使用共享内存存储部分结果
    extern __shared__ float shared_data[];
    shared_data[tid] = 0.0f;
    
    // 每个线程处理多个元素
    int elements_per_thread = (feature_dim + blockDim.x - 1) / blockDim.x;
    int start_idx = tid * elements_per_thread;
    int end_idx = min(start_idx + elements_per_thread, feature_dim);
    
    // 计算分配给这个线程的元素
    for (int i = start_idx; i < end_idx; i++) {
        if (x[base + i] != y[base + i]) {
            shared_data[tid] += 1.0f;
        }
    }
    
    __syncthreads();
    
    // 归约求最终结果
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            shared_data[tid] += shared_data[tid + stride];
        }
        __syncthreads();
    }
    
    // 第一个线程写入结果
    if (tid == 0) {
        distances[sample_idx] = shared_data[0];
    }
}

torch::Tensor hamming_cuda(
    torch::Tensor x,
    torch::Tensor y,
    std::string mode = "fused"
) {
    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") {
        // 融合版本 - 推荐
        const int block_size = 256;
        size_t shared_mem = block_size * sizeof(float);
        
        hamming_kernel_fused<<<batch_size, block_size, shared_mem>>>(
            x_contig.data_ptr<float>(),
            y_contig.data_ptr<float>(),
            distances.data_ptr<float>(),
            batch_size,
            feature_dim
        );
    } else if (mode == "vectorized") {
        // 向量化版本
        const int block_size = 64;  // 每个线程处理4个元素
        size_t shared_mem = block_size * sizeof(float);
        
        hamming_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
        );
    } else {
        // 基础版本
        hamming_kernel_basic<<<batch_size, 1>>>(
            x_contig.data_ptr<float>(),
            y_contig.data_ptr<float>(),
            distances.data_ptr<float>(),
            batch_size,
            feature_dim
        );
    }
    
    return distances;
}
"""

hamming_cpp_source = """
torch::Tensor hamming_cuda(torch::Tensor x, torch::Tensor y, std::string mode);
"""

# Compile the inline CUDA code
hamming = load_inline(
    name="hamming",
    cpp_sources=hamming_cpp_source,
    cuda_sources=hamming_source,
    functions=["hamming_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, mode="fused"):
        super(ModelNew, self).__init__()
        self.mode = mode
        self.hamming = hamming

    def forward(self, x, y):
        return self.hamming.hamming_cuda(x, y, self.mode)
