package cn.edu.bjut.textmining.chapter3;

import cn.edu.bjut.textmining.util.FileUtil;
import cn.edu.bjut.textmining.util.WordListUtil;

import org.deeplearning4j.models.embeddings.learning.impl.elements.CBOW;
import org.deeplearning4j.models.word2vec.VocabWord;
import org.deeplearning4j.models.word2vec.Word2Vec;
import org.deeplearning4j.text.sentenceiterator.BasicLineIterator;
import org.deeplearning4j.text.sentenceiterator.SentenceIterator;
import org.deeplearning4j.text.tokenization.tokenizer.preprocessor.CommonPreprocessor;
import org.deeplearning4j.text.tokenization.tokenizerfactory.DefaultTokenizerFactory;
import org.deeplearning4j.text.tokenization.tokenizerfactory.TokenizerFactory;

import java.nio.file.Path;
import java.util.Collection;
import java.util.List;
import java.util.Set;

/**
 * Word2Vec 词嵌入示例。
 * 语料为 SIGHAN 2005 中文分词评测公开发布的《人民日报》1998 年 1 月语料
 * （北京大学标注版本），语料本身已按词切分。
 * 程序先过滤标点、数字和停用词，得到每行一句的清洗语料，
 * 再用 DeepLearning4J 训练 CBOW 词向量，观察词向量中的语义关系。
 */
public class WordEmbeddingDemo {
    public static void main(String[] args) {
        try {
            Path rawCorpus = FileUtil.path(
                    "data", "chapter3", "sighan2005-pku-training.txt");
            Path stopwordFile = FileUtil.path(
                    "data", "chapter3", "chinese-stopwords-hit.txt");
            Path cleanCorpus = FileUtil.path(
                    "output", "chapter3", "pku-clean.txt");
            Set<String> stopWords = WordListUtil.readLowerCaseSet(stopwordFile);
            int sentenceCount = prepareCorpus(rawCorpus, cleanCorpus, stopWords);

            SentenceIterator iterator = new BasicLineIterator(cleanCorpus.toFile());
            // finish() 释放迭代器持有的文件资源。
            try (AutoCloseable iteratorResource = iterator::finish) {
                TokenizerFactory tokenizerFactory = new DefaultTokenizerFactory();
                tokenizerFactory.setTokenPreProcessor(new CommonPreprocessor());

                Word2Vec vectorModel = new Word2Vec.Builder()
                        .minWordFrequency(5)
                        .epochs(5)
                        .iterations(1)
                        .layerSize(100)
                        .seed(42)
                        .windowSize(5)
                        .workers(1)
                        .negativeSample(10)
                        .elementsLearningAlgorithm(new CBOW<VocabWord>())
                        .iterate(iterator)
                        .tokenizerFactory(tokenizerFactory)
                        .build();

                System.out.println("语料来源: 《人民日报》1998年1月（SIGHAN 2005 公开切分语料）");
                System.out.println("停用词文件: " + stopwordFile);
                System.out.println("清洗后句行数: " + sentenceCount);
                System.out.println("正在训练 Word2Vec 模型...");
                vectorModel.fit();

                String targetWord = "足球";
                if (!vectorModel.hasWord(targetWord)) {
                    throw new IllegalStateException("语料词表中没有“足球”");
                }

                double[] vector = vectorModel.getWordVector(targetWord);
                System.out.println("“足球”的向量前 5 维:");
                for (int i = 0; i < Math.min(5, vector.length); i++) {
                    System.out.printf("%.4f ", vector[i]);
                }
                System.out.println("...");

                Collection<String> nearest = vectorModel.wordsNearest(targetWord, 5);
                System.out.println("与“足球”最相似的词: " + nearest);
                printSimilarity(vectorModel, "足球", "篮球");
                printSimilarity(vectorModel, "足球", "排球");
                printSimilarity(vectorModel, "足球", "经济");
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 语料清洗：原语料已按词切分，逐行过滤非纯汉字词
     * （标点、全角数字、英文字母等）和停用词，写入每行一句的清洗语料。
     */
    private static int prepareCorpus(
            Path source, Path target, Set<String> stopWords) {
        List<String> lines = FileUtil.readUtf8Lines(source);
        StringBuilder builder = new StringBuilder();
        int sentenceCount = 0;
        for (String line : lines) {
            StringBuilder cleaned = new StringBuilder();
            for (String token : line.trim().split("\\s+")) {
                if (!token.matches("[\\u4e00-\\u9fff]+")) {
                    continue;
                }
                if (stopWords.contains(token)) {
                    continue;
                }
                if (cleaned.length() > 0) {
                    cleaned.append(' ');
                }
                cleaned.append(token);
            }
            if (cleaned.length() > 0) {
                builder.append(cleaned).append(System.lineSeparator());
                sentenceCount++;
            }
        }
        FileUtil.writeUtf8Text(target, builder.toString());
        return sentenceCount;
    }

    private static void printSimilarity(
            Word2Vec model, String left, String right) {
        if (model.hasWord(left) && model.hasWord(right)) {
            System.out.printf("“%s”与“%s”的余弦相似度: %.4f%n",
                    left, right, model.similarity(left, right));
        } else {
            System.out.println("跳过相似度计算，词表中缺少: " + right);
        }
    }
}