package com.ruoyi.biemo.nlp;

import com.hankcs.hanlp.corpus.io.IOUtil;
import com.hankcs.hanlp.mining.word2vec.DocVectorModel;
import com.hankcs.hanlp.mining.word2vec.Word2VecTrainer;
import com.hankcs.hanlp.mining.word2vec.WordVectorModel;
import com.ruoyi.biemo.utils.TestUtility;

import java.io.IOException;

/**
 * 词向量工具
 */
public class WordVectorModelUtils {

//    private static final String TRAIN_FILE_NAME = TestUtility.ensureTestData("搜狗文本分类语料库已分词.txt", "http://hanlp.linrunsoft.com/release/corpus/sogou-mini-segmented.zip");
    private static final String TRAIN_FILE_NAME = TestUtility.ensureTestData("搜狗文本分类语料库已分词.txt", "https://develope.lanzoug.com/file/?BGJTbQg5ATAFDFZuADVdMVJtBDwCMQZ8AigHbFF4AXNTOlpuAXEGugPoU6MAslDNVLQOK1/lU/8EigblU4JayARxU2oIPAEjBTZWcgAnXT1SZAQvAjwGdAJuBydRawFpUz9aaAEDBjsDMVM4AGVQYFQ6DjNfNVNkBDMGN1M7Wn0EYlN3CGgBZwVlVjYAZF1lUj8EOQJkBiICdQdxUTABMlNjWj8BbwZ9A2VTMAB7UGBUPw4kX2JTYQRmBmRTOVpsBDNTNwhnATEFYlYwAGddOlI/BGQCZAY2AmAHYFFvAWNTN1o9AW0GNAMzU2cAYVAyVGoOMl8uUzAEcwZrUylaLgR3U2EIJwE8BTdWPQBlXWtSPgQ0Am8GMgI1BydReQFpUz5aaAE6Bm8DZFMzAGxQZ1Q5DjlfOVNoBDYGNVMhWnUEIlNiCDkBIgVuVjEAYl1kUj4ENgJtBjYCNgc0UT0BJlMmWn0BKwZvA2RTMwBsUGJUOQ4yXzFTaQQwBjNTKVouBG1TdAhoAWQFYlYzAHpdblI8BDYCcwY1AjYHN1EjATFTZ1o7");
    private static final String MODEL_FILE_NAME = "/usr/local/textjar/analysis-hanlp/data/test/word2vec.txt";
//    private static final String MODEL_FILE_NAME = "D:\\tianze\\文本大数据\\analysis-hanlp\\data\\test\\word2vec.txt";

    //相似度计算
    public static Float similarCount(String text1,String text2) {
        WordVectorModel wordVectorModel = null;
        try {
            wordVectorModel = trainOrLoadModel();
            DocVectorModel docVectorModel = new DocVectorModel(wordVectorModel);
            return docVectorModel.similarity(text1, text2);
        } catch (IOException ioException) {
            ioException.printStackTrace();
            return null;
        }
    }

    public static void main(String[] args) throws Exception{
        System.out.println(similarCount("沙尔克 VS 斯图加特 初盘 主队 水位 偏高 目前 斯图加特 客场 路有 反弹 迹象 排除 客队 可能 ","沃尔夫斯堡 VS 凯泽斯劳滕 一线 战役 庄家 开出 较高 平局 赔率 想必 有望 分出 胜负"));
    }



    static WordVectorModel trainOrLoadModel() throws IOException
    {
        if (!IOUtil.isFileExisted(MODEL_FILE_NAME))
        {
            if (!IOUtil.isFileExisted(TRAIN_FILE_NAME))
            {
                System.err.println("语料不存在，请阅读文档了解语料获取与格式：https://github.com/hankcs/HanLP/wiki/word2vec");
                System.exit(1);
            }
            Word2VecTrainer trainerBuilder = new Word2VecTrainer();
            return trainerBuilder.train(TRAIN_FILE_NAME, MODEL_FILE_NAME);
        }

        return loadModel();
    }

    static WordVectorModel loadModel() throws IOException
    {
        return new WordVectorModel(MODEL_FILE_NAME);
    }

}
