package cn.edu.bjut.textmining.chapter4;

import weka.classifiers.lazy.IBk;
import weka.core.Attribute;
import weka.core.EuclideanDistance;
import weka.core.Instance;
import weka.core.Instances;
import weka.core.SparseInstance;
import weka.core.neighboursearch.LinearNNSearch;

import java.util.ArrayList;
import java.util.List;

import cn.edu.bjut.textmining.chapter4.support.NewsClassifier;
import cn.edu.bjut.textmining.chapter4.support.NewsRecord;
import cn.edu.bjut.textmining.chapter4.support.TfidfVectorizer;

/** 高维TF-IDF采用线性扫描近邻；关闭逐列极差归一化。 */
public final class WekaNewsClassifier implements NewsClassifier {
    public static final int K = 5;
    private final TfidfVectorizer vectorizer = new TfidfVectorizer();
    private final IBk classifier = new IBk(K);
    private Instances header;
    private double sportsPrior;

    public void fit(List<NewsRecord> training) {
        vectorizer.fit(training);
        int sports = 0;
        for (NewsRecord r : training) {
            if (r.isSports()) {
                sports++;
            }
        }
        sportsPrior = sports / (double) training.size();
        ArrayList<Attribute> attrs = new ArrayList<Attribute>();
        for (String term : vectorizer.terms()) {
            attrs.add(new Attribute(term));
        }
        ArrayList<String> labels = new ArrayList<String>();
        labels.add("finance");
        labels.add("sports");
        attrs.add(new Attribute("class", labels));
        header = new Instances("Tnews", attrs, training.size());
        header.setClassIndex(attrs.size() - 1);
        for (NewsRecord r : training) {
            // 零向量没有方向，不能作为单位TF-IDF的语义近邻。
            if (TfidfVectorizer.isZero(vectorizer.transform(r))) {
                continue;
            }
            Instance x = instance(r);
            x.setValue(header.classIndex(), r.label);
            header.add(x);
        }
        EuclideanDistance distance = new EuclideanDistance();
        distance.setDontNormalize(true);
        LinearNNSearch search = new LinearNNSearch();
        try {
            search.setDistanceFunction(distance);
            classifier.setNearestNeighbourSearchAlgorithm(search);
            classifier.buildClassifier(header);
            System.out.println(
                    "kNN可用训练向量："
                            + header.numInstances()
                            + "；排除零向量："
                            + (training.size() - header.numInstances()));
        } catch (Exception e) {
            throw new IllegalStateException("Weka模型训练失败", e);
        }
    }

    private Instance instance(NewsRecord r) {
        double[] features = vectorizer.transform(r);
        double[] values = new double[features.length + 1];
        System.arraycopy(features, 0, values, 0, features.length);
        values[features.length] = Double.NaN;
        Instance x = new SparseInstance(1.0, values);
        x.setDataset(header);
        return x;
    }

    public double score(NewsRecord r) {
        // 该标题未命中训练词表时，退回训练集类别先验。
        if (TfidfVectorizer.isZero(vectorizer.transform(r))) {
            return sportsPrior;
        }
        try {
            return classifier.distributionForInstance(instance(r))[1];
        } catch (Exception e) {
            throw new IllegalStateException("Weka预测失败", e);
        }
    }

    public double threshold() {
        return 0.5;
    }

    public String name() {
        return "Weka-kNN";
    }
}