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

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

import java.text.Normalizer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Set;

/** 各入口共用同一训练集、测试集和四条观察标题。 */
public final class NewsExperiment {
    public final List<NewsRecord> training;
    public final List<NewsRecord> test;
    public final List<NewsRecord> examples;

    private NewsExperiment(
            List<NewsRecord> training, List<NewsRecord> test, List<NewsRecord> examples) {
        this.training = Collections.unmodifiableList(training);
        this.test = Collections.unmodifiableList(test);
        this.examples = Collections.unmodifiableList(examples);
    }

    public static NewsExperiment load() {
        List<NewsRecord> training = read("train.tsv");
        List<NewsRecord> test = read("test.tsv");
        Set<String> seen = new HashSet<String>();
        for (NewsRecord r : training) {
            if (!seen.add(key(r.text))) {
                throw new IllegalArgumentException("训练集含重复标题");
            }
        }
        for (NewsRecord r : test) {
            if (!seen.add(key(r.text))) {
                throw new IllegalArgumentException("训练/测试集有重复标题");
            }
        }
        List<NewsRecord> examples = new ArrayList<NewsRecord>();
        for (String line : Chapter4Files.read(Chapter4Files.data("examples.tsv"))) {
            if (line.trim().isEmpty()
                    || line.equals("id")
                    || line.startsWith("id\t")
                    || line.startsWith("#")) {
                continue;
            }
            String id = line.split("\t", -1)[0];
            NewsRecord found = null;
            for (NewsRecord r : test) {
                if (r.id.equals(id)) {
                    found = r;
                    break;
                }
            }
            if (found == null) {
                throw new IllegalArgumentException("观察标题不在测试集中：" + id);
            }
            examples.add(found);
        }
        if (examples.size() != 4) {
            throw new IllegalArgumentException("观察标题应为4条");
        }
        return new NewsExperiment(training, test, examples);
    }

    private static List<NewsRecord> read(String name) {
        List<NewsRecord> records = new ArrayList<NewsRecord>();
        for (String line : Chapter4Files.read(Chapter4Files.data(name))) {
            if (line.startsWith("id\t") || line.trim().isEmpty()) {
                continue;
            }
            String[] fields = line.split("\t", -1);
            if (fields.length != 5
                    || !("sports".equals(fields[1]) || "finance".equals(fields[1]))) {
                throw new IllegalArgumentException("数据格式错误：" + name);
            }
            List<String> tokens = ChineseTokenizer.tokenize(fields[2]);
            if (tokens.isEmpty()) {
                throw new IllegalArgumentException("标题分词为空：" + fields[0]);
            }
            records.add(
                    new NewsRecord(
                            fields[0],
                            fields[1],
                            fields[2],
                            fields[3],
                            Integer.parseInt(fields[4]),
                            tokens));
        }
        if (records.isEmpty()) {
            throw new IllegalArgumentException("数据集不能为空：" + name);
        }
        return records;
    }

    private static String key(String text) {
        return Normalizer.normalize(text, Normalizer.Form.NFKC)
                .replaceAll("[\\s\\p{Z}]+", "")
                .toLowerCase(Locale.ROOT);
    }

    public void printSummary() {
        System.out.println("数据：CLUE TNEWS体育/财经新闻标题");
        System.out.println("训练集：" + training.size() + "；测试集：" + test.size());
        System.out.println("体育为正类，财经为负类；四条观察标题均来自测试集。");
    }
}