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

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;

/** 只从训练集学习词表与IDF，随后将任何标题映射到相同坐标系。 */
public final class TfidfVectorizer {
    public static final int MIN_DF = 2;
    public static final int MAX_FEATURES = 2000;
    private final List<String> terms = new ArrayList<String>();
    private final Map<String, Integer> indexes = new HashMap<String, Integer>();
    private double[] idf;

    public void fit(List<NewsRecord> records) {
        final Map<String, Integer> df = new HashMap<String, Integer>();
        for (NewsRecord record : records) {
            for (String token : new HashSet<String>(record.tokens)) {
                df.put(token, df.containsKey(token) ? df.get(token) + 1 : 1);
            }
        }
        terms.clear();
        indexes.clear();
        for (String term : df.keySet()) {
            if (df.get(term) >= MIN_DF) {
                terms.add(term);
            }
        }
        Collections.sort(
                terms,
                new Comparator<String>() {
                    public int compare(String a, String b) {
                        int c = Integer.compare(df.get(b), df.get(a));
                        return c != 0 ? c : a.compareTo(b);
                    }
                });
        if (terms.size() > MAX_FEATURES) {
            terms.subList(MAX_FEATURES, terms.size()).clear();
        }
        if (terms.isEmpty()) {
            throw new IllegalArgumentException("训练词表为空");
        }
        idf = new double[terms.size()];
        for (int i = 0; i < terms.size(); i++) {
            indexes.put(terms.get(i), i);
            idf[i] = Math.log((records.size() + 1.0) / (df.get(terms.get(i)) + 1.0)) + 1.0;
        }
    }

    public double[] transform(NewsRecord record) {
        if (idf == null) {
            throw new IllegalStateException("请先拟合训练集");
        }
        double[] vector = new double[terms.size()];
        for (String token : record.tokens) {
            Integer i = indexes.get(token);
            if (i != null) {
                vector[i]++;
            }
        }
        double squaredNorm = 0.0;
        for (int i = 0; i < vector.length; i++) {
            vector[i] *= idf[i];
            squaredNorm += vector[i] * vector[i];
        }
        if (squaredNorm > 0.0) {
            double norm = Math.sqrt(squaredNorm);
            for (int i = 0; i < vector.length; i++) {
                vector[i] /= norm;
            }
        }
        return vector;
    }

    /** 未命中训练词表时，向量没有可用于距离比较的词项证据。 */
    public static boolean isZero(double[] vector) {
        for (double value : vector) {
            if (value != 0.0) {
                return false;
            }
        }
        return true;
    }

    public List<String> terms() {
        return Collections.unmodifiableList(terms);
    }

    public double idf(int index) {
        return idf[index];
    }
}