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

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

/**
 * @author reghao
 * @date 2025-08-25 11:22:56
 */
public class Bayes {
    public String predictClass(List<List<String>> trainList, List<String> testList) {
        Map<String, List<List<String>>> resultMap = getDataSet(trainList);
        double max = 0.0;
        String result = "";

        for (int i = 0; i < resultMap.size(); i++) {
            double current = 0.0;
            String key = i == 0 ? "yes" : "no";
            List<List<String>> tmpList = resultMap.get(key);
            // 计算总概率 P(y)
            current = DecimalCalculator.divide(tmpList.size(), trainList.size());
            for (int j = 0; j < testList.size(); j++) {
                double pv = calculate(tmpList, testList.get(j), j);
                current = DecimalCalculator.multiply(current, pv);
            }

            if (max <= current) {
                result = i == 0 ? "yes" : "no";
                max = current;
            }
        }

        return result;
    }

    public Map<String, List<List<String>>> getDataSet(List<List<String>> list) {
        Map<String, List<List<String>>> map = new HashMap<>();
        List<List<String>> positiveList = new ArrayList<>();
        List<List<String>> negativeList = new ArrayList<>();
        for (List<String> tmpList : list) {
            // 获取最后一项
            String result = tmpList.get(tmpList.size()-1);
            if (result.equals("yes")) {
                positiveList.add(tmpList);
            } else {
                negativeList.add(tmpList);
            }
        }

        map.put("yes", positiveList);
        map.put("no", negativeList);
        return map;
    }

    public double calculate(List<List<String>> list, String str, int index) {
        int count = 0;
        for (int i = 0; i < list.size(); i++) {
            if (str.equals(list.get(i).get(index))) {
                count++;
            }
        }

        return DecimalCalculator.divide(count, list.size(), 3);
    }
}
