package cn.iocoder.yudao.module.llm.utils;

import com.hankcs.hanlp.HanLP;
import lombok.extern.slf4j.Slf4j;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.ArrayList;
import java.util.List;

/**z
 * @Description : 海明距离算法
 */
@Slf4j
public class HammingUtils {
    public static String getHash (String str) {
        try {
            // 这里使用了MD5获得hash值
            MessageDigest messageDigest = MessageDigest.getInstance("MD5");
            return new BigInteger(1, messageDigest.digest(str.getBytes(StandardCharsets.UTF_8))).toString(2);
        } catch (Exception e) {
            log.error("getHash error:{}", e.getMessage(), e);
            return str;
        }
    }

    /**
     * 传入String,计算出它的simHash值，并以字符串形式输出
     *
     * @param str 传入的String类型字符串
     * @return 返回str的simHash值
     */
    public static String getSimHash(String str) {
        // 用数组表示特征向量,取128位,从 0 1 2 位开始表示从高位到低位
        int[] v = new int[128];
        // 1、分词（使用了外部依赖 hankcs 包提供的接口）
        List<String> keywordList;
        if (str.length() < 200) {
            // 对于短文本，采取不同的处理方式，例如使用更简单的分词或添加额外的处理逻辑
            keywordList = handleShortText(str);
        } else {
            keywordList = HanLP.extractKeyword(str, str.length());// 取出所有关键词
        }
        // hash
        int size = keywordList.size();
        int i = 0; // 以 i 做外层循环
        for (String keyword : keywordList) {
            // 2、获取 hash 值
            String keywordHash = getHash(keyword);
            if (keywordHash.length() < 128) {
                // hash 值可能少于 128 位，在低位以 0 补齐
                int dif = 128 - keywordHash.length();
                for (int j = 0; j < dif; j++) {
                    keywordHash += "0";
                }
            }
            // 3、加权、合并
            for (int j = 0; j < v.length; j++) {
                // 对 keywordHash 的每一位与 '1' 进行比较
                if (keywordHash.charAt(j) == '1') {
                    // 权重分 10 级，由词频从高到低，取权重 10~0
                    v[j] += (10 - (i / (size / 10)));
                } else {
                    v[j] -= (10 - (i / (size / 10)));
                }
            }
            i++;
        }
        // 4、降维
        String simHash = ""; // 储存返回的 simHash 值
        for (int j = 0; j < v.length; j++) {
            // 从高位遍历到低位
            if (v[j] <= 0) {
                simHash += "0";
            } else {
                simHash += "1";
            }
        }
        return simHash;
    }

    private static List<String> handleShortText(String str) {
        // 这里可以添加对短文本的特殊处理逻辑，例如直接使用字符作为关键词
        // 或者使用更简单的分词工具，或者对短文本进行预处理
        // 以下是一个简单的示例，将短文本拆分为单个字符作为关键词
        List<String> result = new ArrayList<>();
        for (char c : str.toCharArray()) {
            result.add(String.valueOf(c));
        }
        return result;
    }

    /** 输入两个 simHash 值，计算它们的海明距离
     *
     * @param simHash1 simHash1
     * @param simHash2 simHash2
     * @return 海明距离
     */
    public static int getHammingDistance(String simHash1, String simHash2) {
        int distance = 0;
        if (simHash1.length() != simHash2.length()) {
            // 出错，返回-1
            distance = -1;
        } else {
            // 将 simHash1 转换为 BigInteger 类型
            BigInteger hash1 = new BigInteger(simHash1, 2);
            // 将 simHash2 转换为 BigInteger 类型
            BigInteger hash2 = new BigInteger(simHash2, 2);
            // 使用 XOR 找出不同的位
            BigInteger xor = hash1.xor(hash2);
            // 计算不同位的数量
            distance = xor.bitCount();
        }
        return distance;
    }

    /**
     * 输入两个 simHash 值，输出相似度
     *
     * @param simHash1 simHash1
     * @param simHash2 simHash2
     * @return 相似度
     */
    public static double getSimilarity(String simHash1, String simHash2) {
        // 通过 simHash1 和 simHash2 获得它们的海明距离
        int distance = getHammingDistance(simHash1, simHash2);
        // 通过海明距离计算出相似度，并返回
        return 0.01 * (100 - (double) (distance * 100) / 128);
    }

}
