import torch
import torch.nn as nn
import torch.nn.functional as F

# -------------------------------------------------------------
# 常量定义
# -------------------------------------------------------------
N_BATCH = 100
D_VECTOR = 128

# 假设 dim=1
DIM = 1


# -------------------------------------------------------------

class Model(nn.Module):
    """
    汉明距离 (Hamming Distance) 的纯 PyTorch 基准实现
    """

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

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        # x1, x2: (N, D), 假设为 long/int/bool 类型

        # 1. 比较 (x1 != x2)
        #    如果 x1=[1, 0, 1], x2=[1, 1, 1]
        #    diff=[False, True, False]
        diff = (x1 != x2)

        # 2. 求和 (Sum)
        #    False.sum() = 0, True.sum() = 1
        #    sum([0, 1, 0]) = 1
        #    我们转换为 float 以匹配 CUDA 版本的输出类型
        return torch.sum(diff, dim=self.dim).to(torch.float32)


def get_inputs():
    """
    生成两个 (N, D) 形状的 *整数* 输入
    (汉明距离的标准输入)
    """
    # 随机生成 0 或 1
    input1 = torch.randint(0, 2, (N_BATCH, D_VECTOR), dtype=torch.long)
    input2 = torch.randint(0, 2, (N_BATCH, D_VECTOR), dtype=torch.long)
    return [input1, input2]


def get_init_inputs():
    return [DIM]