import os
import json


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

    # 数据文件路径
    DATA_DIR = '../Data'
    SCHOOL_COUNT_FILE = os.path.join(DATA_DIR, 'SchoolCount.json')
    TEACHER_COUNT_FILE = os.path.join(DATA_DIR, 'TeacherCount.json')
    ENROLLMENT_COUNT_FILE = os.path.join(DATA_DIR, 'ZaiXiaoShengCount.json')
    ADMISSION_COUNT_FILE = os.path.join(DATA_DIR, 'ZhaoShengCount.json')

    @staticmethod
    def load_json_file(file_path):
        """加载JSON文件数据"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"加载文件 {file_path} 失败: {e}")
            return None

    @staticmethod
    def get_province_data(data):
        """从数据中提取云南省的数据"""
        if not data or not isinstance(data, list):
            return None
        
        # 查找云南省的数据
        for item in data:
            if item.get('area_name') == '云南省':
                return item
        
        print("未找到云南省的数据")
        return None

    @staticmethod
    def get_school_count_by_year_and_stage(year, stage):
        """获取指定年份和学段的学校总数"""
        school_data = EducationDataModel.load_json_file(EducationDataModel.SCHOOL_COUNT_FILE)
        province_data = EducationDataModel.get_province_data(school_data)
        
        if not province_data:
            return 0
        
        # 根据学段获取对应的学校数据字段名
        stage_field_mapping = {
            'preschool': 'preschool_schools',
            'primary': 'primary_schools',
            'junior': 'junior_schools',
            'senior': 'senior_schools',
            'vocational': 'vocational_schools'
        }
        
        field_name = stage_field_mapping.get(stage)
        if not field_name or field_name not in province_data:
            print(f"学段 {stage} 对应的学校数据字段不存在")
            return 0
        
        year_data = province_data[field_name].get(str(year))
        if not year_data:
            print(f"年份 {year} 的学校数据不存在")
            return 0
        
        # 计算学校总数 - 处理vocational_schools的特殊格式
        try:
            # 对于vocational学段，数据可能是直接的数值而不是字典
            if stage == 'vocational' and isinstance(year_data, int):
                return year_data
            # 对于其他学段或正常字典格式
            elif isinstance(year_data, dict):
                if 'urban' in year_data and 'town' in year_data and 'rural' in year_data:
                    total_schools = year_data.get('urban', 0) + year_data.get('town', 0) + year_data.get('rural', 0)
                    return total_schools
                elif 'total' in year_data:
                    # 如果没有分区数据但有total字段，使用total
                    return year_data.get('total', 0)
                else:
                    print(f"年份 {year} 的学校数据缺少必要的字段")
                    return 0
            else:
                print(f"年份 {year} 的学校数据格式不正确")
                return 0
        except Exception as e:
            print(f"计算学校总数时出错: {e}")
            return 0

    @staticmethod
    def get_teacher_count_by_year_and_stage(year, stage):
        """获取指定年份和学段的教职工总数"""
        teacher_data = EducationDataModel.load_json_file(EducationDataModel.TEACHER_COUNT_FILE)
        province_data = EducationDataModel.get_province_data(teacher_data)
        
        if not province_data:
            return 0
        
        # 根据学段获取对应的教职工数据字段名
        stage_field_mapping = {
            'preschool': 'preschool_teachers',
            'primary': 'primary_teachers',
            'junior': 'junior_teachers',
            'senior': 'senior_teachers',
            'vocational': 'vocational_teachers'
        }
        
        field_name = stage_field_mapping.get(stage)
        if not field_name or field_name not in province_data:
            print(f"学段 {stage} 对应的教职工数据字段不存在")
            return 0
        
        year_data = province_data[field_name].get(str(year))
        if not year_data or not isinstance(year_data, dict):
            print(f"年份 {year} 的教职工数据不存在或格式不正确")
            return 0
        
        # 优先使用total_staff，如果为0或不存在，则返回教师总数
        try:
            total_staff = year_data.get('total_staff', 0)
            if total_staff == 0:
                total_staff = year_data.get('total_teacher', 0)
            return total_staff
        except Exception as e:
            print(f"获取教职工总数时出错: {e}")
            return 0

    @staticmethod
    def get_enrollment_count_by_year_and_stage(year, stage):
        """获取指定年份和学段的在校生总数"""
        enrollment_data = EducationDataModel.load_json_file(EducationDataModel.ENROLLMENT_COUNT_FILE)
        province_data = EducationDataModel.get_province_data(enrollment_data)
        
        if not province_data or 'student_data' not in province_data:
            print("未找到学生数据")
            return 0
        
        stage_data = province_data['student_data'].get(stage)
        if not stage_data:
            print(f"学段 {stage} 的学生数据不存在")
            return 0
        
        year_data = stage_data.get(str(year))
        if not year_data or not isinstance(year_data, dict):
            print(f"年份 {year} 的学生数据不存在或格式不正确")
            return 0
        
        # 返回在校生总数 - 处理vocational数据只有total字段的情况
        try:
            if 'urban' in year_data and 'town' in year_data and 'rural' in year_data:
                # 通过累加城区、镇区、乡村数据计算总数
                total_enrollment = year_data.get('urban', 0) + year_data.get('town', 0) + year_data.get('rural', 0)
                return total_enrollment
            elif 'total' in year_data and stage == 'vocational':
                # 对于vocational学段，如果只有total字段，使用total
                return year_data.get('total', 0)
            else:
                # 对于其他学段，如果没有分区数据，直接报错
                print(f"年份 {year} 的学生数据缺少城区、镇区或乡村字段")
                return 0
        except Exception as e:
            print(f"获取在校生总数时出错: {e}")
            return 0

    @staticmethod
    def get_admission_count_by_year_and_stage(year, stage):
        """获取指定年份和学段的招生总数"""
        admission_data = EducationDataModel.load_json_file(EducationDataModel.ADMISSION_COUNT_FILE)
        province_data = EducationDataModel.get_province_data(admission_data)
        
        if not province_data or 'education_data' not in province_data:
            print("未找到教育数据")
            return 0
        
        stage_data = province_data['education_data'].get(stage)
        if not stage_data:
            print(f"学段 {stage} 的教育数据不存在")
            return 0
        
        year_data = stage_data.get(str(year))
        if not year_data or not isinstance(year_data, dict):
            print(f"年份 {year} 的教育数据不存在或格式不正确")
            return 0
        
        # 返回招生总数 - 处理vocational数据只有total字段的情况
        try:
            if 'urban' in year_data and 'town' in year_data and 'rural' in year_data:
                # 通过累加城区、镇区、乡村数据计算总数
                total_admission = year_data.get('urban', 0) + year_data.get('town', 0) + year_data.get('rural', 0)
                return total_admission
            elif 'total' in year_data and stage == 'vocational':
                # 对于vocational学段，如果只有total字段，使用total
                return year_data.get('total', 0)
            else:
                # 对于其他学段，如果没有分区数据，直接报错
                print(f"年份 {year} 的教育数据缺少城区、镇区或乡村字段")
                return 0
        except Exception as e:
            print(f"获取招生总数时出错: {e}")
            return 0

    @staticmethod
    def get_education_data_by_year(year):
        """获取指定年份所有学段的教育数据"""
        result = []
        
        for stage_code, stage_name in EducationDataModel.EDUCATION_STAGES.items():
            # 获取各类数据（增加异常处理）
            try:
                school_count = EducationDataModel.get_school_count_by_year_and_stage(year, stage_code)
                teacher_count = EducationDataModel.get_teacher_count_by_year_and_stage(year, stage_code)
                enrollment_count = EducationDataModel.get_enrollment_count_by_year_and_stage(year, stage_code)
                admission_count = EducationDataModel.get_admission_count_by_year_and_stage(year, stage_code)
                
                # 转换为万人单位
                teacher_count_10k = round(teacher_count / 10000, 2) if teacher_count else 0
                enrollment_count_10k = round(enrollment_count / 10000, 2) if enrollment_count else 0
                admission_count_10k = round(admission_count / 10000, 2) if admission_count else 0
                
                # 添加到结果列表
                result.append({
                    'education_stage': stage_name,
                    'school_count': school_count,
                    'teacher_count_10k': teacher_count_10k,
                    'enrollment_count_10k': enrollment_count_10k,
                    'admission_count_10k': admission_count_10k
                })
            except Exception as e:
                print(f"处理学段 {stage_name} 时出错: {e}")
                # 出错时添加一个默认的空数据条目
                result.append({
                    'education_stage': stage_name,
                    'school_count': 0,
                    'teacher_count_10k': 0,
                    'enrollment_count_10k': 0,
                    'admission_count_10k': 0
                })
        
        return result

    @staticmethod
    def print_education_data(data, year):
        """打印教育数据"""
        print(f"\n===== {year}年云南省各学段教育数据统计（单位：学校总数-个，其余-万人）=====")
        print(f"{'学段':<8}{'学校总数':<10}{'教职工总数':<12}{'在校生总数':<12}{'招生总数':<12}")
        print("=" * 60)
        
        for item in data:
            print(f"{item['education_stage']:<8}{item['school_count']:<10}{item['teacher_count_10k']:<12.2f}{item['enrollment_count_10k']:<12.2f}{item['admission_count_10k']:<12.2f}")

    @staticmethod
    def main(year=2023):
        """主函数"""
        try:
            print(f"正在获取 {year} 年数据...")
            print(f"数据目录: {EducationDataModel.DATA_DIR}")
            
            # 检查数据文件是否存在
            for file_path in [EducationDataModel.SCHOOL_COUNT_FILE, 
                             EducationDataModel.TEACHER_COUNT_FILE, 
                             EducationDataModel.ENROLLMENT_COUNT_FILE, 
                             EducationDataModel.ADMISSION_COUNT_FILE]:
                if os.path.exists(file_path):
                    print(f"找到数据文件: {file_path}")
                else:
                    print(f"警告：数据文件不存在: {file_path}")
            
            # 获取数据
            education_data = EducationDataModel.get_education_data_by_year(year)
            
            # 打印数据
            if education_data:
                EducationDataModel.print_education_data(education_data, year)
            else:
                print("未获取到任何数据")
                
        except Exception as e:
            print(f"获取数据时发生错误: {e}")
            import traceback
            traceback.print_exc()


if __name__ == "__main__":
    EducationDataModel.main()