package cn.edu.bjut.textmining.chapter7;

import cc.mallet.pipe.CharSequence2TokenSequence;
import cc.mallet.pipe.Pipe;
import cc.mallet.pipe.SerialPipes;
import cc.mallet.pipe.TokenSequence2FeatureSequence;
import cc.mallet.pipe.iterator.ArrayIterator;
import cc.mallet.topics.ParallelTopicModel;
import cc.mallet.topics.TopicInferencer;
import cc.mallet.types.IDSorter;
import cc.mallet.types.InstanceList;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.TreeSet;
import java.util.regex.Pattern;
import java.util.Locale;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;

/**
 * 使用 MALLET 训练一个小型 LDA，并输出主题词与文档主题分布。
 */
public class LdaMalletExample {

    public static void main(String[] args) {
        try {
            List<String> documents = Chapter7Data.readNonBlankLines("lda_documents.txt");
            ArrayList<Pipe> pipes = new ArrayList<Pipe>();
            pipes.add(new CharSequence2TokenSequence(Pattern.compile("\\S+")));
            pipes.add(new TokenSequence2FeatureSequence());
            InstanceList instances = new InstanceList(new SerialPipes(pipes));
            instances.addThruPipe(new ArrayIterator(documents));
            int topicCount = 3;
            ParallelTopicModel model = new ParallelTopicModel(topicCount, 0.2, 0.01);
            model.setRandomSeed(2026);
            model.addInstances(instances);
            model.setNumThreads(1);
            model.setNumIterations(500);
            model.setOptimizeInterval(0);
            model.setTopicDisplay(0, 0);
            model.estimate();
            System.out.println("实验性质：30 行人工组织的词项序列，用于观察主题分布；不是新闻分类评测。");
            System.out.println("主题编号由模型产生，不能直接解释为体育、科技或经济类别。");
            System.out.println("--- 每个主题的关键词 ---");
            ArrayList<TreeSet<IDSorter>> topicWords = model.getSortedWords();
            for (int topic = 0; topic < topicCount; topic++) {
                System.out.print("主题" + topic + "：");
                int count = 0;
                for (IDSorter sorter : topicWords.get(topic)) {
                    System.out.print(instances.getDataAlphabet().lookupObject(sorter.getID()) + " ");
                    count++;
                    if (count >= 8) {
                        break;
                    }
                }
                System.out.println();
            }
            System.out.println("\n--- 从输入词项读取主题比例（各组展示两条）---");
            int[] preview = { 0, 1, 10, 11, 20, 21 };
            for (int d : preview) {
                System.out.println("doc-" + d + "｜输入：" + documents.get(d));
                System.out.println("  " + TopicOutput.distribution(model.getTopicProbabilities(d)));
            }
            List<String> rows = new ArrayList<String>();
            rows.add("document_id\ttokens\ttopic_0\ttopic_1\ttopic_2");
            for (int d = 0; d < instances.size(); d++) {
                double[] weights = model.getTopicProbabilities(d);
                rows.add(String.format(Locale.ROOT, "doc-%d\t%s\t%.8f\t%.8f\t%.8f", d, documents.get(d), weights[0], weights[1], weights[2]));
            }
            Files.write(Chapter7Data.outputPath("lda_document_topics.tsv"), rows, StandardCharsets.UTF_8);
            System.out.println("完整 30 条文档及比例：output/chapter7/lda_document_topics.tsv");
            System.out.println("比例表示一篇文档的主题混合程度，四舍五入后各项之和可能略偏离 1。");
            // 与训练输入相同，使用空格分隔的词项。
            String newDoc = args.length == 0 ? "大模型 训练 推理 芯片 算力" : String.join(" ", args);
            // 新文档只能使用训练时确定的词表，不在推断时扩展词编号。
            instances.getDataAlphabet().stopGrowth();
            String knownWords = retainKnownWords(newDoc, instances);
            TopicInferencer inferencer = model.getInferencer();
            inferencer.setRandomSeed(2026);
            InstanceList test = new InstanceList(instances.getPipe());
            test.addThruPipe(new ArrayIterator(Arrays.asList(knownWords)));
            double[] inferred = inferencer.getSampledDistribution(test.get(0), 200, 20, 20);
            System.out.println("\n新文档: " + newDoc);
            System.out.println("新文档主题分布：" + TopicOutput.distribution(inferred));
            System.out.println("推断固定训练主题与词表，仅估计新文档的主题比例；该比例不是分类置信度。");
        } catch (Exception e) {
            throw new IllegalStateException("运行失败，请检查输入文件和依赖。", e);
        }
    }

    static String retainKnownWords(String document, InstanceList training) {
        List<String> known = new ArrayList<String>();
        List<String> unknown = new ArrayList<String>();
        String normalized = document.trim();
        if (!normalized.isEmpty()) {
            for (String token : normalized.split("\\s+")) {
                if (training.getDataAlphabet().lookupIndex(token, false) >= 0) {
                    known.add(token);
                } else {
                    unknown.add(token);
                }
            }
        }
        int total = known.size() + unknown.size();
        System.out.printf(Locale.ROOT, "新文档词表覆盖：%d/%d 个词项实例%n", known.size(), total);
        System.out.println("未登录词：" + unknown);
        if (known.isEmpty()) {
            throw new IllegalArgumentException("新文档没有训练词表内的有效词项，不能据纯先验分布解释其主题。");
        }
        return String.join(" ", known);
    }
}