package com.ruoyi.biemo.nlp;

import com.hankcs.hanlp.classification.classifiers.IClassifier;
import com.hankcs.hanlp.classification.classifiers.NaiveBayesClassifier;
import com.hankcs.hanlp.classification.models.NaiveBayesModel;
import com.hankcs.hanlp.corpus.io.IOUtil;
import com.ruoyi.biemo.utils.TestUtility;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;

public class  TextClassificationUtils {

    /**
     * 搜狗文本分类语料库5个类目，每个类目下1000篇文章，共计5000篇文章
     */
    public static final String CORPUS_FOLDER = TestUtility.ensureTest2Data("搜狗文本分类语料库迷你版", "http://hanlp.linrunsoft.com/release/corpus/sogou-text-classification-corpus-mini.zip");
    /**
     * 模型保存路径
     */
    public static final String MODEL_PATH = "/usr/local/textjar/analysis-hanlp/data/test/classification-model.ser";
//    public static final String MODEL_PATH = "D:\\tianze\\文本大数据\\analysis-hanlp\\data\\test\\classification-model.ser";

    //文本分类
    public static String getClassification(String text){
        String result = "";
        try {
            IClassifier classifier = new NaiveBayesClassifier(trainOrLoadModel());
            result =  classifier.classify(text);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return result;

    }
    //训练模型
    private static NaiveBayesModel trainOrLoadModel() throws IOException
    {
        NaiveBayesModel model = (NaiveBayesModel) IOUtil.readObjectFrom(MODEL_PATH);
        if (model != null) return model;

        File corpusFolder = new File(CORPUS_FOLDER);
        if (!corpusFolder.exists() || !corpusFolder.isDirectory())
        {
            System.err.println("没有文本分类语料，请阅读IClassifier.train(java.lang.String)中定义的语料格式与语料下载：") ;
            System.exit(1);
        }

        IClassifier classifier = new NaiveBayesClassifier(); // 创建分类器，更高级的功能请参考IClassifier的接口定义
        classifier.train(CORPUS_FOLDER);                     // 训练后的模型支持持久化，下次就不必训练了
        model = (NaiveBayesModel) classifier.getModel();
        IOUtil.saveObjectTo(model, MODEL_PATH);
        return model;
    }

    //我的训练模型
    public static Object myTrainOrLoadModel(MultipartFile[] files, String name) {
        if(files!=null&&files.length>0){
            File file1 = new File(CORPUS_FOLDER+"/"+name);
            if(!file1.getParentFile().exists()){
                file1.mkdirs();
            }
            if(!file1.exists()){
                file1.mkdirs();
            }
            for(int i=0;i<files.length;i++){
                MultipartFile file = files[i];
                try {
                    File file2 =  new File(file1.getCanonicalPath()+"/"+ file.getOriginalFilename());
                    file.transferTo(file2);
                } catch (IOException ioException) {
                    ioException.printStackTrace();
                }
            }
            IClassifier classifier = new NaiveBayesClassifier();
            try {
                classifier.train(CORPUS_FOLDER);
                NaiveBayesModel model = (NaiveBayesModel) classifier.getModel();
                IOUtil.saveObjectTo(model, MODEL_PATH);
            } catch (IOException ioException) {
                ioException.printStackTrace();
            }
        }
        return "ok";
    }

}
