package cn.edu.bjut.textmining.chapter5.support;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;

import cn.edu.bjut.textmining.util.TextVectorizer;

/** 字符N-gram TF-IDF；所有文档共用一个词表和同一组DF。 */
public final class TextFeatures {
    private TextFeatures() {}

    private static final Set<String> STOP_GRAMS =
            new HashSet<>(
                    Arrays.asList(
                            "什么", "怎么", "如何", "为什么", "一个", "一种", "可以", "哪些", "有没有", "是不是", "这个",
                            "这样", "中国", "现在", "我们", "他们", "你们", "自己", "的人", "的是", "知道", "多少", "进行",
                            "问题", "时候", "为何", "来看", "认为"));

    public static final class Result {
        public final double[][] vectors;
        public final List<String> vocabulary;
        public final int[] documentFrequencies;

        Result(double[][] vectors, List<String> vocabulary, int[] df) {
            this.vectors = vectors;
            this.vocabulary = vocabulary;
            this.documentFrequencies = df;
        }
    }

    public static List<String> ngrams(String text, int minN, int maxN, boolean filter) {
        // 标点、符号和空白划定片段边界，不把两侧字符拼成原文中不存在的组合。
        String clean =
                java.text.Normalizer.normalize(text, java.text.Normalizer.Form.NFKC)
                        .toLowerCase(Locale.ROOT);
        List<String> result = new ArrayList<>();
        for (String segment : clean.split("[\\p{P}\\p{S}\\s]+")) {
            int[] codePoints = segment.codePoints().toArray();
            for (int n = minN; n <= maxN; n++) {
                for (int i = 0; i + n <= codePoints.length; i++) {
                    String gram = new String(codePoints, i, n);
                    if (!filter || (!STOP_GRAMS.contains(gram) && !gram.matches("[0-9]+"))) {
                        result.add(gram);
                    }
                }
            }
        }
        return result;
    }

    public static Result fit(List<String> documents, boolean largeCorpus) {
        if (documents.isEmpty()) {
            throw new IllegalArgumentException("语料不能为空");
        }
        Map<String, Integer> frequency = new HashMap<>();
        List<List<String>> tokenized = new ArrayList<>();
        for (String document : documents) {
            List<String> grams = ngrams(document, largeCorpus ? 2 : 1, largeCorpus ? 3 : 2, largeCorpus);
            tokenized.add(grams);
            for (String gram : new HashSet<>(grams)) {
                frequency.put(gram, frequency.getOrDefault(gram, 0) + 1);
            }
        }
        List<String> vocabulary = new ArrayList<>(frequency.keySet());
        if (largeCorpus) {
            vocabulary.removeIf(
                    gram ->
                            frequency.get(gram) < 2
                                    || frequency.get(gram) > documents.size() * 0.60);
            vocabulary.sort(
                    (a, b) -> {
                        int order = frequency.get(b).compareTo(frequency.get(a));
                        return order != 0 ? order : a.compareTo(b);
                    });
            if (vocabulary.size() > 5000) {
                vocabulary = new ArrayList<>(vocabulary.subList(0, 5000));
            }
        } else {
            Collections.sort(vocabulary);
        }
        Map<String, Integer> index = new HashMap<>();
        int[] df = new int[vocabulary.size()];
        for (int j = 0; j < vocabulary.size(); j++) {
            index.put(vocabulary.get(j), j);
            df[j] = frequency.get(vocabulary.get(j));
        }
        double[][] vectors = new double[documents.size()][vocabulary.size()];
        for (int i = 0; i < documents.size(); i++) {
            for (String gram : tokenized.get(i)) {
                Integer j = index.get(gram);
                if (j != null) {
                    vectors[i][j]++;
                }
            }
            for (int j = 0; j < vocabulary.size(); j++) {
                double idf = Math.log((documents.size() + 1.0) / (df[j] + 1.0)) + 1.0;
                vectors[i][j] *= idf;
            }
            TextVectorizer.normalize(vectors[i]);
        }
        return new Result(vectors, vocabulary, df);
    }
}