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

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

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Locale;

/** 混淆矩阵与按平均秩计算的ROC AUC；不要求分数是概率。 */
public final class EvaluationSupport {
    private EvaluationSupport() {}

    public static final class Scored {
        public final NewsRecord record;
        public final double score;

        public Scored(NewsRecord record, double score) {
            if (!Double.isFinite(score)) {
                throw new IllegalArgumentException("预测分数须为有限数");
            }
            this.record = record;
            this.score = score;
        }
    }

    public static final class Metrics {
        public int tp, fp, fn, tn;
        public double accuracy, precision, recall, f1, auc;

        public String row(String name) {
            return String.format(
                    Locale.ROOT,
                    "%s\t%d\t%d\t%d\t%d\t%.6f\t%.6f\t%.6f\t%.6f\t%.6f\n",
                    name,
                    tp,
                    fp,
                    fn,
                    tn,
                    accuracy,
                    precision,
                    recall,
                    f1,
                    auc);
        }
    }

    public static List<Scored> score(NewsClassifier model, List<NewsRecord> test) {
        List<Scored> results = new ArrayList<Scored>();
        for (NewsRecord r : test) {
            results.add(new Scored(r, model.score(r)));
        }
        return results;
    }

    public static Metrics calculate(List<Scored> scores, double threshold) {
        if (scores.isEmpty()) {
            throw new IllegalArgumentException("测试集不能为空");
        }
        Metrics m = new Metrics();
        for (Scored s : scores) {
            if (s.record.isSports()) {
                if (s.score >= threshold) {
                    m.tp++;
                } else {
                    m.fn++;
                }
            } else {
                if (s.score >= threshold) {
                    m.fp++;
                } else {
                    m.tn++;
                }
            }
        }
        m.accuracy = (m.tp + m.tn) / (double) scores.size();
        m.precision = divide(m.tp, m.tp + m.fp);
        m.recall = divide(m.tp, m.tp + m.fn);
        m.f1 = divide(2.0 * m.precision * m.recall, m.precision + m.recall);
        m.auc = auc(scores);
        return m;
    }

    private static double divide(double a, double b) {
        return b == 0.0 ? 0.0 : a / b;
    }

    private static double auc(List<Scored> scores) {
        List<Scored> sorted = new ArrayList<Scored>(scores);
        Collections.sort(
                sorted,
                new Comparator<Scored>() {
                    public int compare(Scored a, Scored b) {
                        return Double.compare(a.score, b.score);
                    }
                });
        int positive = 0, negative = 0, i = 0;
        double rankSum = 0;
        while (i < sorted.size()) {
            int end = i + 1;
            while (end < sorted.size()
                    && Double.compare(sorted.get(i).score, sorted.get(end).score) == 0) {
                end++;
            }
            double rank = (i + 1 + end) / 2.0;
            for (int j = i; j < end; j++) {
                if (sorted.get(j).record.isSports()) {
                    positive++;
                    rankSum += rank;
                } else {
                    negative++;
                }
            }
            i = end;
        }
        if (positive == 0 || negative == 0) {
            return Double.NaN;
        }
        return (rankSum - positive * (positive + 1) / 2.0) / (positive * (double) negative);
    }

    public static void saveScores(String name, List<Scored> scores, double threshold) {
        StringBuilder out = new StringBuilder("id\tlabel\tscore\tprediction\ttext\n");
        for (Scored s : scores) {
            out.append(s.record.id)
                    .append('\t')
                    .append(s.record.label)
                    .append('\t')
                    .append(s.score)
                    .append('\t')
                    .append(s.score >= threshold ? "sports" : "finance")
                    .append('\t')
                    .append(s.record.text)
                    .append('\n');
        }
        Chapter4Files.write(Chapter4Files.output("predictions-" + name + ".tsv"), out.toString());
    }

    public static Metrics evaluate(NewsClassifier model, NewsExperiment e) {
        List<Scored> scored = score(model, e.test);
        Metrics m = calculate(scored, model.threshold());
        System.out.printf(
                Locale.ROOT,
                "%s: Accuracy=%.2f%%, Precision=%.2f%%, Recall=%.2f%%, F1=%.2f%%, AUC=%.4f%n",
                model.name(),
                100 * m.accuracy,
                100 * m.precision,
                100 * m.recall,
                100 * m.f1,
                m.auc);
        System.out.printf("TP=%d, FP=%d, FN=%d, TN=%d%n", m.tp, m.fp, m.fn, m.tn);
        saveScores(model.name(), scored, model.threshold());
        return m;
    }

    public static void showExamples(NewsClassifier model, NewsExperiment e) {
        for (int i = 0; i < e.examples.size(); i++) {
            NewsRecord r = e.examples.get(i);
            double score = model.score(r);
            System.out.println("N" + (i + 1) + " " + r.id + "：" + r.text);
            System.out.println("分词：" + r.tokens);
            System.out.printf(
                    Locale.ROOT,
                    "真实=%s，预测=%s，%s=%.6f%n",
                    r.chineseLabel(),
                    score >= model.threshold() ? "体育" : "财经",
                    model instanceof LibSvmNewsClassifier ? "体育方向决策值" : "体育类别得分",
                    score);
        }
    }
}