#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
import os, json
from Initial_Data_Program.Initial_Data_Calculate_std_and_mean import calculate_pic_std_and_mean
from Initial_Data_Program.Initial_Data_Gen_configs_base_datasets_my_dataset import generate_configs_base_datasets_my_dataset_file
from Initial_Data_Program.Initial_Data_Gen_mmseg_utils_class_names import generate_mmseg_utils_class_names_file
from Initial_Data_Program.Initial_Data_Gen_mmseg_datasets_my_dataset import generate_mmseg_datasets_my_dataset_file
from Initial_Data_Program.Initial_Data_Gen_mmseg_datasets_init_ import generate_mmseg_datasets_init_file

def load_json_files(directory, not_check_list=None):
    """
    读取指定文件夹下的所有 JSON 文件，排除指定的文件。
    
    :param directory: 要读取的目录
    :param not_check_list: 要排除的文件名列表，不包含路径。默认为空列表
    :return: 包含所有有效 JSON 数据的列表，每个元素为一个字典
    """
    if not_check_list is None:
        not_check_list = ['']  # 默认排除文件

    # 获取所有 .json 文件，并排除在 not_check_list 中的文件
    json_files = [f for f in os.listdir(directory) if f.endswith('.json') and f not in not_check_list]
    
    data_list = []
    
    # 遍历每个 JSON 文件
    for json_file in json_files:
        json_path = os.path.join(directory, json_file)
        
        try:
            # 打开并加载 JSON 文件
            with open(json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                # 将文件名（去掉 .json）添加到数据中
                data['file_name_json'] = json_file.rstrip(".json")
                data_list.append(data)
        
        except json.JSONDecodeError:
            print(f"\033[91mError decoding JSON file: {json_path}\033[0m")
        except Exception as e:
            print(f"\033[91mError reading file {json_path}: {str(e)}\033[0m")

    return data_list


def process_json_data(json_data):
    """处理每个 JSON 数据，生成相应参数"""
    # 提取 dataset 信息
    dataset_file_name = json_data["dataset_info"]["dataset_file_name"]
    dataset_class_name = json_data["dataset_info"]["dataset_class_name"]
    data_root = json_data["dataset_info"]["data_root"]
    
    # 转换 img_scale 为元组 (img_scale_width, img_scale_height)
    img_scale = (json_data["dataset_info"]["img_scale_width"], json_data["dataset_info"]["img_scale_height"])
    
    # 提取其他必要信息
    train_img_path = json_data["dataset_info"]["paths"]["train_img_path"]
    train_seg_map_path = json_data["dataset_info"]["paths"]["train_seg_map_path"]
    val_img_path = json_data["dataset_info"]["paths"]["val_img_path"]
    val_seg_map_path = json_data["dataset_info"]["paths"]["val_seg_map_path"]
    test_img_path = json_data["dataset_info"]["paths"]["test_img_path"]
    test_seg_map_path = json_data["dataset_info"]["paths"]["test_seg_map_path"]
    
    # 提取 label 相关信息
    classes = json_data["label_info"]["classes"]
    palette = json_data["label_info"]["palette"]
    img_suffix = json_data["label_info"]["img_suffix"]
    seg_map_suffix = json_data["label_info"]["seg_map_suffix"]
    reduce_zero_label = json_data["label_info"]["reduce_zero_label"]

    # 提取训练相关参数
    # 转换 crop_size 为元组 (crop_size_width, crop_size_height)
    crop_size = (json_data["training_info"]["crop_size_width"], json_data["training_info"]["crop_size_height"])
    train_batch_size = json_data["training_info"]["train_batch_size"]
    train_num_workers = json_data["training_info"]["train_num_workers"]
    val_and_test_batch_size = json_data["training_info"]["val_and_test_batch_size"]
    val_and_test_num_workers = json_data["training_info"]["val_and_test_num_workers"]

    return (dataset_file_name, dataset_class_name, data_root, img_scale, train_img_path, train_seg_map_path, val_img_path, val_seg_map_path, test_img_path, test_seg_map_path,
            classes, palette, img_suffix, seg_map_suffix, reduce_zero_label, crop_size, train_batch_size, train_num_workers, val_and_test_batch_size, val_and_test_num_workers,)


def save_all_record_to_json(output_file, dataset_file_names, classes_all, palette_all, palette_num_all, mean_all, std_all):
    """构建一个 JSON 文件，用于存储每个数据集的信息"""
    # 构建一个字典，用于存储每个数据集的信息
    data_record = {}

    # 假设所有列表长度相同，遍历每个 dataset_file_name
    for i in range(len(dataset_file_names)):
        data_record[dataset_file_names[i]] = {
            "classes": classes_all[i],
            "palette": palette_all[i],
            "palette_num": palette_num_all[i],
            "mean": list(mean_all[i]),
            "std": list(std_all[i])
        }

    # 将字典写入 JSON 文件
    with open(output_file, 'w', encoding='utf-8') as json_file:
        json.dump(data_record, json_file, ensure_ascii=False, indent=6)

    print(f"\033[93m相关数据汇总到 {output_file} successfully!\033[0m")
    return True

if __name__ == '__main__':
    ########### 1.超参数-定义训练文件夹路径 ###########
    train_parameter_dir = './My_All_In_One/1_Data_Parameter'
    all_data_record_json = "All_Data_Record.json" # 记录所有数据

    ########### 2.1.遍历所有配置文件，并生成对应数据集 ###########
    # 定义保存信息的列表
    dataset_file_names = []
    dataset_class_names = []
    classes_all = []
    palette_all = []
    palette_num_all = []
    mean_all = []
    std_all = []

    # 2.1.1. 从 ./1_Data_Parameter 文件夹读取所有 JSON 文件
    json_data_list = load_json_files(train_parameter_dir, not_check_list = [all_data_record_json])

    # 2.1.2. 遍历每个 JSON 数据文件并生成对应配置文件
    for json_data in json_data_list:
        # A. 输出当前处理文件
        print(f"\033[32m正在处理{json_data['file_name_json']}.json文件\033[0m")
        
        # B. 处理 JSON 数据并提取参数
        (dataset_file_name, dataset_class_name, data_root, img_scale, train_img_path, train_seg_map_path, val_img_path, val_seg_map_path, test_img_path, test_seg_map_path,
         classes, palette, img_suffix, seg_map_suffix, reduce_zero_label, crop_size, train_batch_size, train_num_workers, val_and_test_batch_size, val_and_test_num_workers,
         ) = process_json_data(json_data)

        # 保存文件名和类别名
        dataset_file_names.append(dataset_file_name)
        dataset_class_names.append(dataset_class_name)

        # C. 文件存储位置
        output_configs_base_datasets_my_dataset = f'./configs/_base_/datasets/{dataset_file_name}.py'
        output_mmseg_datasets_dataset_file_name = os.path.join(f'./mmseg/datasets/{dataset_file_name}.py')
        output_mmseg_datasets_init = os.path.join('./mmseg/datasets/__init__.py')
        output_mmseg_utils_class_names = f'./mmseg/utils/class_names.py'

        # D. 运行程序生成配置文件
        # 生成 ./configs/_base_/datasets/{dataset_file_name}.py
        print("    ",end='')
        success = generate_configs_base_datasets_my_dataset_file(
            output_file=output_configs_base_datasets_my_dataset,
            dataset_class_name=dataset_class_name,
            data_root=data_root,
            img_scale=img_scale,
            crop_size=crop_size,
            train_batch_size=train_batch_size,
            train_num_workers=train_num_workers,
            val_and_test_batch_size=val_and_test_batch_size,
            val_and_test_num_workers=val_and_test_num_workers,
            train_img_path=train_img_path,
            train_seg_map_path=train_seg_map_path,
            val_img_path=val_img_path,
            val_seg_map_path=val_seg_map_path,
            test_img_path=test_img_path,
            test_seg_map_path=test_seg_map_path
        )

        # 生成 ./mmseg/datasets/{dataset_file_name}.py
        print("    ",end='')
        success, classes, palette = generate_mmseg_datasets_my_dataset_file(
            output_file=output_mmseg_datasets_dataset_file_name,
            dataset_class_name=dataset_class_name,
            classes=classes,
            palette=palette,
            img_suffix=img_suffix,
            seg_map_suffix=seg_map_suffix,
            reduce_zero_label=reduce_zero_label
        )

        mean, std = calculate_pic_std_and_mean(
            dataset_dir = os.path.join(data_root, train_img_path)
        )

        # 保存标注类名和颜色
        classes_all.append(classes)
        palette_all.append(palette)
        palette_num_all.append(len(classes))
        mean_all.append(mean)
        std_all.append(std)

    ########### 2.2.汇总所有信息运行生成 init 和 class_names 文件 ###########
    # 生成 ./mmseg/datasets/__init__.py
    print("    ",end='')
    success = generate_mmseg_datasets_init_file(
        output_file=output_mmseg_datasets_init,
        dataset_file_names=dataset_file_names,
        dataset_class_names=dataset_class_names
    )

    # 生成 ./mmseg/utils/class_names.py
    print("    ",end='')
    success = generate_mmseg_utils_class_names_file(
        output_file=output_mmseg_utils_class_names,
        dataset_file_names=dataset_file_names,
        classes_all=classes_all,
        palette_all=palette_all
    )

    ########### 2.2.汇总dataset_file_names、classes_all、palette_all、palette_num_all所有信息到My_All_In_One/1_Data_Parameter/All_Data_Record.json文件 ###########
    output_all_data_record = os.path.join(train_parameter_dir, all_data_record_json)
    # 调用函数保存数据
    success = save_all_record_to_json(output_file=output_all_data_record, dataset_file_names=dataset_file_names, classes_all=classes_all, palette_all=palette_all, palette_num_all=palette_num_all, mean_all=mean_all, std_all=std_all)
    