package cn.reghao.tnb.sb.bayes.email;

import lombok.Getter;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author reghao
 * @date 2025-08-25 16:37:51
 */
@Getter
public class SpamTrain {
    TextProcessor textProcessor;
    Map<String, Integer> spamMailMap = new HashMap<>();// 垃圾邮件中的词语 -> 每个词语出现的次数
    Map<String, Integer> hamMailMap = new HashMap<>();// 正常邮件中的词语 -> 每个词语出现的次数
    Integer spamMailSegNum = 0;// 垃圾邮件中词语的总数量
    Integer hamMailSegNum = 0;// 正常邮件中词语的总数量
    Map<String, Double> spamMailRateMap = new HashMap<>(); // 垃圾邮件中每个词语出现的概率
    Map<String, Double> hamMailRateMap = new HashMap<>();// 正常邮件中每个词语出现的概率
    int trainNum = 0; // 训练的邮件数量
    Integer spamTrainMailNum = 0; // 垃圾邮件的数量
    double spamTrainRate = 0d; // 垃圾邮件率
    double trainTime = 0d; // 训练时间

    public SpamTrain(TextProcessor textProcessor) {
        this.textProcessor = textProcessor;
    }

    void train(String trainDataPath) {
        long trainStartTime = System.currentTimeMillis();
        List<String> list = textProcessor.getTextLines(trainDataPath);
        for (String line : list) {
            trainNum++;
            String[] arr = line.split(" ");
            String type = arr[0];
            String contentPath = arr[1];
            System.out.println(contentPath);
            List<String> wordList = textProcessor.getWordList(contentPath);

            // 根据邮件类型分别处理
            if (type.equals("spam")) {
                spamTrainMailNum ++;

                spamMailSegNum += wordList.size();
                for (String word : wordList) {
                    spamMailMap.put(word, spamMailMap.containsKey(word) ? spamMailMap.get(word) + 1 : 1);
                }
            } else {
                hamMailSegNum += wordList.size();
                for (String word : wordList) {
                    hamMailMap.put(word, hamMailMap.containsKey(word) ? hamMailMap.get(word) + 1 : 1);
                }
            }
        }

        // 垃圾邮件率
        spamTrainRate = (double) spamTrainMailNum / (double) trainNum;
        setSpamMailRateMap();

        long trainFinishTime = System.currentTimeMillis();
        trainTime = (trainFinishTime - trainStartTime) / 1000D;
    }

    void setSpamMailRateMap() {
        for (Map.Entry<String, Integer> entry : spamMailMap.entrySet()) {
            String word = entry.getKey();
            int occurCount = entry.getValue();

            // P1(ti) 某个词语在垃圾邮件的词语中出现的概率
            double rate = (double) occurCount / (double) spamMailSegNum;
            // 某个词语在正常邮件中出现的次数, 若没有出现过则默认为 1 次(防止判断过于绝对)
            int occurCountHam = hamMailMap.getOrDefault(word, 1);
            // P2(ti) 某个词语在正常邮件的词语中出现的概率
            double rateHam = (double) occurCountHam/(double) hamMailSegNum;
            // P(ti) 某个词语在所有邮件的词语中出现的概率
            // P(ti) = P1(ti) + P2(ti)
            double allRate = rate + rateHam;
            // 某个词语在邮件中出现时, 该邮件为垃圾邮件的概率
            // P1(ti)/P(ti) = P1(ti)/(P1(ti) + P2(ti)) = P(spam|ti)
            //
            // 某个词语在邮件中出现时, 该邮件为正常邮件的概率
            // P2(ti)/P(ti) = P2(ti)/(P1(ti) + P2(ti)) = P(ham|ti)
            double spamRate = rate/allRate;
            // 词语 -> 垃圾邮件概率
            spamMailRateMap.put(word, spamRate);
        }
    }
}
