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

import lombok.Getter;

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

/**
 * @author reghao
 * @date 2025-08-25 16:43:39
 */
@Getter
public class SpamTest {
    TextProcessor textProcessor;
    Map<String, Integer> spamMailMap;
    Map<String, Integer> hamMailMap;
    Integer spamMailSegNum;
    Integer hamMailSegNum;
    Map<String, Double> spamMailRateMap;
    int trainNum;
    Integer spamTrainMailNum;
    double spamTrainRate;

    double[] probability; // 测试邮件为垃圾邮件的概率数组
    String[] testRealType; // 测试邮件真实的类型数组
    int testNum = 0; // 实际测试邮件数
    double testTime = 0D; // 测试时间

    public SpamTest(SpamTrain spamTrain, int maxTestNum) {
        this.textProcessor = spamTrain.getTextProcessor();
        this.spamMailMap = spamTrain.getSpamMailMap();
        this.hamMailMap = spamTrain.getHamMailMap();
        this.spamMailSegNum = spamTrain.getSpamMailSegNum();
        this.hamMailSegNum = spamTrain.getHamMailSegNum();
        this.spamMailRateMap = spamTrain.getSpamMailRateMap();
        this.trainNum = spamTrain.getTrainNum();
        this.spamTrainMailNum = spamTrain.getSpamTrainMailNum();
        this.spamTrainRate = spamTrain.getSpamTrainRate();
        this.probability = new double[maxTestNum];
        this.testRealType = new String[maxTestNum];
    }

    void test(String testDataPath) {
        long testStartTime = System.currentTimeMillis();
        List<String> list = textProcessor.getTextLines(testDataPath);
        for (String line : list) {
            // 每次处理一封邮件
            String[] arr = line.split(" ");
            String type = arr[0];
            String contentPath = arr[1];
            List<String> wordList = textProcessor.getWordList(contentPath);

            // 标识当前邮件是否手动设置概率
            int flag = 0;
            // P(spam) 是所有邮件中垃圾邮件的概率(先验概率)
            // P(spam) = spamTrainRate
            // P(spam) * {P(spam|ti) * P(ti) / P(spam)} (i 从 1 ~ n)
            double rate = spamTrainRate;
            // P(ham) = 1 - P(spam) = 1 - spamTrainRate
            // (1 - P(spam)) * {(1 - P(spam|ti)) * P(ti) / P(spam)} (i 从 1 ~ n)
            double tmpRate = 1 - spamTrainRate;
            for (String word : wordList) {
                if (spamMailRateMap.containsKey(word)) {
                    // P(spam|ti)
                    double spamTmp = spamMailRateMap.get(word);

                    int occurCount = spamMailMap.get(word) + hamMailMap.getOrDefault(word, 0);
                    int segNum = spamMailSegNum + hamMailSegNum;
                    // P(ti)
                    double rate0 = (double) occurCount / (double) segNum;

                    // P(spam) * {P(spam|ti) * P(ti) / P(spam)}
                    rate *= spamTmp * rate0 / spamTrainRate;
                    // (1 - P(spam)) * {(1 - P(spam|ti)) * P(ti) / P(spam)}
                    tmpRate *= (1D - spamTmp) * rate0 / spamTrainRate;
                    if (testNum == 1 || testNum == 2) {
                        System.out.println("rate:"+rate+"\ntempRate:"+tmpRate);
                    }

                    // 当某个概率非常趋近 0 时需要判断该邮件更接近垃圾邮件还是正常邮件
                    if (rate < Math.pow(10, -300) || tmpRate < Math.pow(10, -300)) {
                        if (rate > 1) {
                            probability[testNum] = 1D;
                            flag = 2;
                            break;
                        } else if (tmpRate > 1) {
                            probability[testNum] = 0D;
                            flag = 3;
                            break;
                        } else {
                            rate *= Math.pow(10, 7);
                            tmpRate *= Math.pow(10, 7);
                        }
                    }
                }
            }

            if (flag == 0) {
                // flag != 0 表示已经手动设置过概率
                probability[testNum] = rate / (rate + tmpRate);
                System.out.println(probability[testNum]);
            }

            testRealType[testNum++] = type;
        }

        long testFinishTime = System.currentTimeMillis();
        testTime = (testFinishTime - testStartTime) / 1000D;
    }
}
