import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
import math

N_BATCH = 100
D_VECTOR = 128

DIM = 1
BLOCK_SIZE = 256


class ModelNew(nn.Module):

    def __init__(self, dim=1):
        super().__init__()

        if dim != 1:
            raise NotImplementedError("CUDA Kernel only supports dim=1 for (N,D) tensors")
        self.dim = dim
        self.block_size = BLOCK_SIZE
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_header = f"""
        #include <torch/extension.h>

        // C++ 接口
        torch::Tensor hamming_distance_forward_cuda(
            torch::Tensor x1,
            torch::Tensor x2,
            int dim
        );
        """

        cuda_source = f"""
        #include <torch/extension.h> 
        #include <cuda_runtime.h>
        #include <cmath> 

        #define BLOCK_SIZE {self.block_size}

        /*
         * Hamming Distance 融合核函数
         * (输入: int64_t, 输出: float)
         */
        __global__ void hamming_distance_fused_kernel(
            const int64_t* __restrict__ x1_data,
            const int64_t* __restrict__ x2_data,
            float* __restrict__ output_data,
            int N, 
            int D
        ) {{

            __shared__ int s_data[BLOCK_SIZE];

            const int n_idx = blockIdx.x; // [0, N-1]
            const int tid = threadIdx.x;  // [0, BLOCK_SIZE-1]

            const int64_t* p_in1 = x1_data + (int64_t)n_idx * D;
            const int64_t* p_in2 = x2_data + (int64_t)n_idx * D;


            int thread_sum = 0;

            for (int d = tid; d < D; d += BLOCK_SIZE) {{
                if (p_in1[d] != p_in2[d]) {{
                    thread_sum++;
                }}
            }}
            s_data[tid] = thread_sum;
            __syncthreads();


            for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {{
                if (tid < offset) {{
                    s_data[tid] += s_data[tid + offset];
                }}
                __syncthreads();
            }}


            if (tid == 0) {{
                output_data[n_idx] = (float)s_data[0];
            }}
        }}

        // C++ 封装函数
        torch::Tensor hamming_distance_forward_cuda(
            torch::Tensor x1,
            torch::Tensor x2,
            int dim
        ) {{

            TORCH_CHECK(x1.dim() == 2, "CUDA Kernel only supports 2D (N, D) input");
            TORCH_CHECK(dim == 1, "CUDA Kernel only supports dim=1");
            TORCH_CHECK(x1.is_cuda() && x2.is_cuda(), "Inputs must be CUDA");
            TORCH_CHECK(x1.is_contiguous() && x2.is_contiguous(), "Inputs must be contiguous");
            TORCH_CHECK(x1.sizes() == x2.sizes(), "Input shapes must match");


            TORCH_CHECK(x1.scalar_type() == torch::kLong, "Input 1 must be torch.long");
            TORCH_CHECK(x2.scalar_type() == torch::kLong, "Input 2 must be torch.long");

            const int64_t N = x1.size(0);
            const int64_t D = x1.size(1);


            auto output_options = torch::TensorOptions()
                                    .dtype(torch::kFloat32)
                                    .device(x1.device());
            auto output = torch::empty({{N}}, output_options);

            dim3 grid_dim(N);
            dim3 block_dim(BLOCK_SIZE);

            hamming_distance_fused_kernel<<<grid_dim, block_dim>>>(
                x1.data_ptr<int64_t>(),
                x2.data_ptr<int64_t>(),
                output.data_ptr<float>(),
                N, D
            );

            return output;
        }}
        """

        nvcc_flags = ['-O3']

        # JIT (Just-In-Time) 编译
        self.hamming_op = load_inline(
            name="hamming_dist_op_v1",
            cpp_sources=cpp_header,
            cuda_sources=cuda_source,
            functions=["hamming_distance_forward_cuda"],
            extra_cuda_cflags=nvcc_flags,
            verbose=False
        )

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        # 确保连续性
        x1_cont = x1.contiguous()
        x2_cont = x2.contiguous()

        # CUDA 核函数要求 Long 类型
        # (如果输入不是 Long, .to() 会创建一个副本)
        x1_long = x1_cont.to(torch.long)
        x2_long = x2_cont.to(torch.long)

        return self.hamming_op.hamming_distance_forward_cuda(
            x1_long,
            x2_long,
            self.dim
        )