package cn.edu.bjut.textmining.chapter4;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import cn.edu.bjut.textmining.chapter4.support.NewsClassifier;
import cn.edu.bjut.textmining.chapter4.support.NewsRecord;
import cn.edu.bjut.textmining.chapter4.support.TfidfVectorizer;
import opennlp.tools.doccat.DoccatFactory;
import opennlp.tools.doccat.DoccatModel;
import opennlp.tools.doccat.DocumentCategorizerME;
import opennlp.tools.doccat.DocumentSample;
import opennlp.tools.ml.AbstractTrainer;
import opennlp.tools.ml.naivebayes.NaiveBayesTrainer;
import opennlp.tools.util.CollectionObjectStream;
import opennlp.tools.util.ObjectStream;
import opennlp.tools.util.TrainingParameters;

/** OpenNLP朴素贝叶斯；重复词元保留，训练词表由训练数据建立。 */
public final class OpenNlpNewsClassifier implements NewsClassifier {
    private DocumentCategorizerME categorizer;
    private int sportsIndex;
    private final Set<String> vocabulary = new HashSet<String>();
    private double sportsPrior;

    public void fit(List<NewsRecord> training) {
        TfidfVectorizer vectorizer = new TfidfVectorizer();
        vectorizer.fit(training);
        vocabulary.clear();
        vocabulary.addAll(vectorizer.terms());
        int sports = 0;
        for (NewsRecord r : training) {
            if (r.isSports()) {
                sports++;
            }
        }
        sportsPrior = sports / (double) training.size();
        List<DocumentSample> samples = new ArrayList<DocumentSample>();
        for (NewsRecord r : training) {
            samples.add(new DocumentSample(r.label, filteredTokens(r)));
        }
        TrainingParameters p = new TrainingParameters();
        p.put(TrainingParameters.CUTOFF_PARAM, 1);
        p.put(AbstractTrainer.ALGORITHM_PARAM, NaiveBayesTrainer.NAIVE_BAYES_VALUE);
        try (ObjectStream<DocumentSample> stream =
                new CollectionObjectStream<DocumentSample>(samples)) {
            DoccatModel model = DocumentCategorizerME.train("zh", stream, p, new DoccatFactory());
            categorizer = new DocumentCategorizerME(model);
            sportsIndex = categorizer.getIndex("sports");
        } catch (Exception e) {
            throw new IllegalStateException("OpenNLP模型训练失败", e);
        }
    }

    public double score(NewsRecord r) {
        if (categorizer == null) {
            throw new IllegalStateException("请先训练模型");
        }
        String[] tokens = filteredTokens(r);
        return tokens.length == 0 ? sportsPrior : categorizer.categorize(tokens)[sportsIndex];
    }

    private String[] filteredTokens(NewsRecord record) {
        List<String> words = new ArrayList<String>();
        for (String word : record.tokens) {
            if (vocabulary.contains(word)) {
                words.add(word);
            }
        }
        return words.toArray(new String[0]);
    }

    public double threshold() {
        return 0.5;
    }

    public String name() {
        return "OpenNLP-NB";
    }
}