import json


class RuYuanZaiYuanModel:
    # 定义支持的教育阶段映射
    EDUCATION_STAGES = {
        'preschool': '学前',
        'primary': '小学',
        'junior': '初中',
        'senior': '高中'
    }

    @staticmethod
    def load_student_data():
        try:
            # 加载招生数据(入园人数)
            with open("./Data/ZhaoShengCount.json", "r", encoding="utf-8") as f:
                enrollment_data = json.load(f)

            # 加载在校生数据(在园人数)
            with open("./Data/ZaiXiaoShengCount.json", "r", encoding="utf-8") as f:
                in_school_data = json.load(f)

            return enrollment_data, in_school_data
        except Exception as e:
            print(f"读取学生数据出错: {e}")
            return [], []

    @staticmethod
    def generate_preschool_education_config(education_stage='preschool'):
        # 验证教育阶段参数
        if education_stage not in RuYuanZaiYuanModel.EDUCATION_STAGES:
            education_stage = 'preschool'  # 默认使用学前

        # 获取学前教育相关数据
        enrollment_data, in_school_data = RuYuanZaiYuanModel.load_student_data()
        # 提取云南省级数据
        yunnan_enroll = next((item for item in enrollment_data if item["area_name"] == "云南省"), None)

        if not yunnan_enroll:
            return {}

        # # 构建学前教育数据
        urban_data = []  # 城区数据
        town_data = []  # 镇区数据
        rural_data = []  # 乡村数据
        total_enroll = []  # 总人数

        # 提取年份数据(2015-2024)
        years = [str(year) for year in range(2015, 2025)]

        for year in years:
            # 使用传入的教育阶段参数
            enroll_data = yunnan_enroll["education_data"].get(education_stage, {}).get(year, {})
            urban_data.append(enroll_data.get("urban", 0) / 10000)  # 转换为万人
            town_data.append(enroll_data.get("town", 0) / 10000)  # 转换为万人
            rural_data.append(enroll_data.get("rural", 0) / 10000)  # 转换为万人
            # 计算总和作为总人数
            calculated_total = enroll_data.get("urban", 0) + enroll_data.get("town", 0) + enroll_data.get("rural", 0)
            total_enroll.append(calculated_total / 10000)  # 转换为万人

        # 添加2022年基数的粉色折线
        base_year = "2022"
        # 找到2022年在years中的索引位置
        base_index = years.index(base_year) if base_year in years else 0
        # 获取2022年的总人数作为基数
        base_value = total_enroll[base_index] if base_index < len(total_enroll) else 0
        # 创建2022年基数折线数据（2022-2024年）
        base_2022_line = []
        for i, year in enumerate(years):
            # 只在2022年及之后显示基数线
            if i >= base_index:
                base_2022_line.append(base_value)
            else:
                base_2022_line.append(None)  # 2022年之前不显示
        data = {
            "xAxis_data": years,
            "series_data_0": urban_data,  # 城区
            "series_data_1": town_data,  # 镇区
            "series_data_2": rural_data,  # 乡村
            "series_data_3": total_enroll,  # 总人数
            "series_data_4": base_2022_line,  # 2022年基数
            "education_stage": RuYuanZaiYuanModel.EDUCATION_STAGES.get(education_stage, '学前')  # 添加教育阶段名称
        }
        return data

    @staticmethod
    def generate_in_school_education_config(education_stage='preschool'):
        # 验证教育阶段参数
        if education_stage not in RuYuanZaiYuanModel.EDUCATION_STAGES:
            education_stage = 'preschool'  # 默认使用学前

        # 获取在校生相关数据
        enrollment_data, in_school_data = RuYuanZaiYuanModel.load_student_data()

        # 提取云南省级数据
        yunnan_in_school = next((item for item in in_school_data if item["area_name"] == "云南省"), None)

        if not yunnan_in_school:
            return {}

        # 提取年份数据(2015-2024)
        years = [str(year) for year in range(2015, 2025)]

        # 构建数据
        urban_data = []  # 城区数据
        town_data = []  # 镇区数据
        rural_data = []  # 乡村数据
        total_in_school = []  # 总人数

        for year in years:
            # 使用传入的教育阶段参数
            in_school_year_data = yunnan_in_school["student_data"].get(education_stage, {}).get(year, {})

            # 先获取原始数据
            urban_val = in_school_year_data.get("urban", 0)
            town_val = in_school_year_data.get("town", 0)
            rural_val = in_school_year_data.get("rural", 0)

            # 转换为万人并添加到各自列表
            urban_data.append(urban_val / 10000)  # 转换为万人
            town_data.append(town_val / 10000)  # 转换为万人
            rural_data.append(rural_val / 10000)  # 转换为万人

            # 计算总和并转换为万人
            calculated_total = (urban_val + town_val + rural_val) / 10000  # 先计算总和再转换为万人
            total_in_school.append(calculated_total)

        # 添加2022年基数的粉色折线
        base_year = "2022"
        # 找到2022年在years中的索引位置
        base_index = years.index(base_year) if base_year in years else 0
        # 获取2022年的总人数作为基数
        base_value = total_in_school[base_index] if base_index < len(total_in_school) else 0
        # 创建2022年基数折线数据（2022-2024年）
        base_2022_line = []
        for i, year in enumerate(years):
            # 只在2022年及之后显示基数线
            if i >= base_index:
                base_2022_line.append(base_value)
            else:
                base_2022_line.append(None)  # 2022年之前不显示

        data = {
            "xAxis_data": years,
            "series_data_0": urban_data,
            "series_data_1": town_data,
            "series_data_2": rural_data,
            "series_data_3": total_in_school,
            "series_data_4": base_2022_line,
            "education_stage": RuYuanZaiYuanModel.EDUCATION_STAGES.get(education_stage, '学前')  # 添加教育阶段名称
        }

        return data
