import os, json, subprocess
from Initial_Schedule_Program.Initial_Train_Gen_configs_base_schedules_schedule_XXk import generate_times_configs_base_schedules_schedule_file
from Initial_Schedule_Program.Initial_Train_Gen_configs_base_schedules_schedule_XXe import generate_epochs_configs_base_schedules_schedule_file
from datetime import datetime

from Initial_Alg_Program.Initial_Alg_Gen_Tool import format_all_data, write_config_to_file, get_gpu_info
from Initial_Alg_Program.Initial_Alg_Gen_0_base_ import generate_base_config
from Initial_Alg_Program.Initial_Alg_Gen_1_norm_cfg import generate_norm_cfg
from Initial_Alg_Program.Initial_Alg_Gen_2_data_preprocessor import generate_data_preprocessor
from Initial_Alg_Program.Initial_Alg_Gen_3_model import generate_model_pretrained, generate_model_backbone, generate_model_data_preprocessor, generate_model_decode_head, generate_model_auxiliary_head, generate_model_train_cfg, generate_model_test_cfg

def load_json_file(file_name):
    """
    读取并返回指定 JSON 文件的内容。

    :param file_name: 要读取的 JSON 文件路径
    :return: JSON 文件的内容作为字典或列表
    """
    try:
        # 打开并读取 JSON 文件
        with open(file_name, 'r', encoding='utf-8') as f:
            data = json.load(f)
        return data
    except FileNotFoundError:
        print(f"\033[91mError: File {file_name} not found.\033[0m")
        return None
    except json.JSONDecodeError:
        print(f"\033[91mError: Failed to decode JSON from {file_name}.\033[0m")
        return None
    except Exception as e:
        print(f"\033[91mAn error occurred while reading {file_name}: {str(e)}\033[0m")
        return None

def process_data_record_json_data(json_data):
    """
    将 JSON 数据直接存入一个大的字典 all_data_record。
    :param json_data: 从 JSON 文件加载的数据
    :return: all_data_record 字典
    """
    all_data_record = {}

    # 遍历 json_data 中的每个键（即每个数据集）
    for dataset_name, dataset_info in json_data.items():
        all_data_record[dataset_name] = {
            "classes": dataset_info['classes'],
            "palette": dataset_info['palette'],
            "palette_num": dataset_info['palette_num'],
            "mean": dataset_info['mean'],
            "std": dataset_info['std'],
            "train_imgs_num": dataset_info['train_imgs_num']
        }
    
    return all_data_record

# 选择要处理的数据集
def select_dataset(all_data_record):
    """
    让用户从 all_data_record 中选择数据集，并返回 dataset_file_name 和对应的 palette_num。
    
    :param all_data_record: 包含所有数据集信息的字典
    :return: 选定的 dataset_file_name 和 palette_num
    """
    # 获取所有数据集的名称
    dataset_names = list(all_data_record.keys())

    # 显示可用的数据集，并让用户选择
    print("选择可用数据集：")
    for i, name in enumerate(dataset_names):
        print(f"{i + 1}. {name} - {all_data_record[name]['palette_num']}类")
    
    # 用户输入选择的数据集编号
    while True:
        try:
            selection = int(input("请选择数据集编号（输入数字）：")) - 1
            if 0 <= selection < len(dataset_names):
                dataset_file_name = dataset_names[selection]
                break
            else:
                print(f"输入的编号无效，请输入 1 到 {len(dataset_names)} 之间的数字。")
        except ValueError:
            print("无效输入，请输入数字。")

    # 获取对应的 palette_num
    palette_num = all_data_record[dataset_file_name]['palette_num']
    mean = all_data_record[dataset_file_name]['mean']
    std = all_data_record[dataset_file_name]['std']
    train_imgs_num = all_data_record[dataset_file_name]['train_imgs_num']
    
    print(f"    已选择数据集: {dataset_file_name} ，其对应的分类数为: {palette_num}")
    
    return dataset_file_name, palette_num, mean, std, train_imgs_num

# 选择对应的算法
def select_alg(alg_directory):
    """
    选择一个算法文件并运行它。
    
    :return: 选定的算法文件名和相对路径
    """
    # 获取 ./Alg 目录下的所有 Python 文件
    algorithms = [f for f in os.listdir(alg_directory) if f.endswith('.py')]
    
    # 显示可用的算法文件
    print("选择可用算法：")
    for i, alg in enumerate(algorithms):
        print(f"{i + 1}. {alg}")
    
    # 用户选择算法
    while True:
        try:
            selection = int(input("请选择算法编号（输入数字）：")) - 1
            if 0 <= selection < len(algorithms):
                selected_alg = algorithms[selection]
                break
            else:
                print(f"输入的编号无效，请输入 1 到 {len(algorithms)} 之间的数字。")
        except ValueError:
            print("无效输入，请输入数字。")
    
    # 生成选定算法的相对路径
    relative_alg_path = os.path.join(alg_directory, selected_alg)
    print(f"    已选择算法: {selected_alg}  {relative_alg_path}")
    
    return selected_alg, relative_alg_path

# 选择计算用的GPU
def select_GPU():
    """
    让用户选择要使用的 GPU 数量，并返回相应的 GPU 列表（以 0,1 格式输出）。
    如果用户没有输入任何选择，默认为选择一块 GPU，编号为 0。
    :return: 用户选择的 GPU 列表和数量
    """
    # 获取 GPU 信息
    gpu_info = get_gpu_info()
    
    if not gpu_info:
        return None, 0
    
    # 显示 GPU 信息
    print("可用的 GPU 列表：")
    for idx, mem_free in gpu_info:
        print(f"GPU {idx}: 剩余显存 {mem_free} MB")
    
    # 提供默认选择 GPU 个数和编号
    default_num_gpus = 1
    default_gpu_idx = 0
    
    # 用户选择 GPU 个数（默认选择 1）
    try:
        num_gpus = input(f"请选择使用的 GPU 个数 (1-{len(gpu_info)}, 默认为 1): ")
        num_gpus = int(num_gpus) if num_gpus else default_num_gpus
        if not (1 <= num_gpus <= len(gpu_info)):
            print(f"输入无效，使用默认值 1 个 GPU。")
            num_gpus = default_num_gpus
    except ValueError:
        print(f"无效输入，使用默认值 1 个 GPU。")
        num_gpus = default_num_gpus
    
    # 用户选择 GPU 编号（如果选择多个 GPU，逐个选择编号）
    selected_gpus = []
    for i in range(num_gpus):
        try:
            gpu_idx = input(f"请输入要使用的第 {i + 1} 个 GPU 的编号 (默认为 {default_gpu_idx}): ")
            gpu_idx = int(gpu_idx) if gpu_idx else default_gpu_idx
            if gpu_idx in [idx for idx, _ in gpu_info] and gpu_idx not in selected_gpus:
                selected_gpus.append(gpu_idx)
            else:
                print(f"无效输入，使用默认 GPU {default_gpu_idx}")
                selected_gpus.append(default_gpu_idx)
        except ValueError:
            print(f"无效输入，使用默认 GPU {default_gpu_idx}")
            selected_gpus.append(default_gpu_idx)
    
    # 返回以 0,1 格式的 GPU 列表和 GPU 数量
    gpu_list_str = ','.join(map(str, selected_gpus))
    print(f"\n已选择 GPU: {gpu_list_str}，共 {num_gpus} 块 GPU")
    
    return gpu_list_str, num_gpus

# 选择训练批次相关信息
def select_schedule(train_imgs_num):
    """
    让用户选择训练模式（Iteration 或 Epoch），并根据选择收集参数，
    最终生成对应的训练计划 schedule 配置文件。
    """
    # 1. 让用户选择模式
    while True:
        mode = input("请选择训练模式 (1: Iteration, 2: Epoch) [默认: 2]: ").strip()
        if mode in ['1', '2', '']:
            mode = '2' if mode == '' else mode
            break
        else:
            print("无效输入，请输入 1 或 2。")

    # --- 模式 1: 基于 Iteration 的配置 ---
    if mode == '1':
        print("\n--- 您已选择 Iteration 模式 ---")
        try:
            train_time_or_epoch_k = input("请输入训练次数 (k为单位, 默认40k): ").strip()
            train_time_or_epoch_k = int(train_time_or_epoch_k) if train_time_or_epoch_k else 40
        except ValueError:
            print("无效输入，使用默认训练时间 40k")
            train_time_or_epoch_k = 40

        try:
            check_num = input("请输入总的验证/保存次数 (默认10): ").strip()
            check_num = int(check_num) if check_num else 10
        except ValueError:
            print("无效输入，使用默认次数 10")
            check_num = 10

        try:
            loggerhook_interval = input("请输入日志间隔 (默认50次迭代): ").strip()
            loggerhook_interval = int(loggerhook_interval) if loggerhook_interval else 50
        except ValueError:
            print("无效输入，使用默认日志间隔 50")
            loggerhook_interval = 50

        val_proportion = 1 / check_num
        # 根据验证比例计算间隔，确保至少为1
        interval = max(1, int(train_time_or_epoch_k * 1000 * val_proportion))

        schedule_file_name = f'schedule_{train_time_or_epoch_k}k_check_{interval}.py'
        output_path = os.path.join('./configs/_base_/schedules/', schedule_file_name)

        generate_times_configs_base_schedules_schedule_file(
            output_file=output_path,
            train_time_or_epoch_k=train_time_or_epoch_k,
            val_proportion=val_proportion,
            loggerhook_interval=loggerhook_interval
        )
        # 返回文件名和训练时长
        return schedule_file_name, "Iteration", train_time_or_epoch_k

    # --- 模式 2: 基于 Epoch 的配置 ---
    elif mode == '2':
        print("\n--- 您已选择 Epoch 模式 ---")
        try:
            max_epochs = input("请输入训练总轮数 (Epochs, 默认300): ").strip()
            max_epochs = int(max_epochs) if max_epochs else 300
        except ValueError:
            print("无效输入，使用默认轮数 300")
            max_epochs = 300

        try:
            val_interval = input("请输入验证间隔的轮次数 (Epochs, 默认1): ").strip()
            val_interval = int(val_interval) if val_interval else 1
        except ValueError:
            print("无效输入，使用默认验证间隔 1")
            val_interval = 1

        try:
            checkpoint_interval = input("请输入模型保存间隔的轮次数 (Epochs, 默认10): ").strip()
            checkpoint_interval = int(checkpoint_interval) if checkpoint_interval else 10
        except ValueError:
            print("无效输入，使用默认保存间隔 10")
            checkpoint_interval = 10
        
        try:
            loggerhook_interval_default = train_imgs_num // 16
            loggerhook_interval = input(f"请输入日志间隔 (默认{loggerhook_interval_default}次迭代): ").strip()
            loggerhook_interval = int(loggerhook_interval) if loggerhook_interval else loggerhook_interval_default
        except ValueError:
            print(f"无效输入，使用默认日志间隔 {loggerhook_interval_default}")
            loggerhook_interval = loggerhook_interval_default

        schedule_file_name = f'schedule_{max_epochs}e_val{val_interval}_check{checkpoint_interval}.py'
        output_path = os.path.join('./configs/_base_/schedules/', schedule_file_name)

        # 注意：这里调用的是基于Epoch的生成函数
        generate_epochs_configs_base_schedules_schedule_file(
            output_file=output_path,
            max_epochs=max_epochs,
            val_interval=val_interval,
            checkpoint_interval=checkpoint_interval,
            loggerhook_interval=loggerhook_interval
        )
        # 返回文件名和训练时长
        return schedule_file_name, "Epoch", max_epochs

# 运行对应的代码去生成算法配置文件
def run_alg_to_gen_alg(relative_alg_path, mean, std, alg_name, dataset_file_name, num_classes, GPU_num, schedule_file_name, train_type, train_time_or_epoch_k):
    try:
        # 构建命令和参数
        cmd = [
            'python', relative_alg_path,
            '--alg_name', alg_name,
            '--dataset_file_name', dataset_file_name,
            '--mean', str(mean),
            '--std', str(std),
            '--dataset_num_classes', str(num_classes),
            '--GPU_num', str(GPU_num),
            '--schedule_file_name', schedule_file_name,
            '--train_type', train_type,
            '--train_time_or_epoch_k', str(train_time_or_epoch_k)
        ]
        
        # 打印当前执行的命令
        print(f"\n正在运行: \033[33m{' '.join(cmd)}\033[0m")
        
        # 运行 Python 脚本并传递参数
        subprocess.run(cmd, check=True)
        
        print(f"\033[92m算法生成器 {relative_alg_path} 运行成功。\033[0m")
    except subprocess.CalledProcessError as e:
        print(f"\033[91m算法生成器 {relative_alg_path} 运行失败，错误信息: {e}\033[0m")

if __name__ == '__main__':
    ########### 1.超参数-定义数据集、算法信息路径 ###########
    train_parameter_dir = './My_All_In_One/1_Data_Parameter'
    all_data_record_json = "All_Data_Record.json"
    all_data_record_file = os.path.join(train_parameter_dir, all_data_record_json) # 数据集信息记录路径
    
    alg_directory = './My_All_In_One/2_Alg_Program' # 算法配置生成路径

    work_dir_base = '../DataSet_Public_outputs' # 工作路径

    ########### 2.获取现有数据集信息 ###########
    print(f"\033[36m{'='*10} 一、选择训练数据集 {'='*10}\033[0m")
    data_record_json_data = load_json_file(file_name = all_data_record_file) # 加载数据集信息
    all_data_record = process_data_record_json_data(json_data = data_record_json_data) # 分析数据集信息
    dataset_file_name, num_classes, mean, std, train_imgs_num = select_dataset(all_data_record = all_data_record) # 选择特定数据集

    ########### 3.获取现有GPU信息 ###########
    print(f"\033[36m{'='*10} 二、选择训练GPU {'='*10}\033[0m")
    gpu_list_str, GPU_num = select_GPU() # 选择GPU

    ########### 4.获取训练批次信息 ###########
    print(f"\033[36m{'='*10} 三、选择训练批次信息 {'='*10}\033[0m")
    schedule_file_name, train_type, train_time_or_epoch_k = select_schedule(train_imgs_num) # 选择训练批次
    schedule_file_name = schedule_file_name.rstrip('.py') 

    ########### 5.获取现有算法信息 ###########
    print(f"\033[36m{'='*10} 四、选择训练算法 {'='*10}\033[0m")
    selected_alg, relative_alg_path = select_alg(alg_directory=alg_directory) # 选择算法
    alg_name = selected_alg.rstrip(".py")
    
    ########### 6.运行选定的算法 ###########
    print(f"\033[36m{'='*5} 生成训练算法配置 {'='*5}\033[0m")
    run_alg_to_gen_alg(relative_alg_path=relative_alg_path, alg_name=alg_name, dataset_file_name=dataset_file_name, num_classes=num_classes, GPU_num=GPU_num, mean=mean, std=std,schedule_file_name=schedule_file_name, train_type = train_type, train_time_or_epoch_k=train_time_or_epoch_k)

    # 算法相关信息
    with open("_temp_.txt", "r") as file:
        data = json.load(file)
        alg_file_name = data['alg_infos']['alg_file_name']
        alg_file_pth = data['alg_infos']['alg_file_pth']
    os.remove("_temp_.txt")

    ########### 7.输出工作、算法目录 ###########
    # 获取并打印当前的年月日时分秒
    data_now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    work_dir_file = os.path.join(work_dir_base, f"{dataset_file_name}-Class_{num_classes}-Alg_{alg_name}-AlgName_{alg_file_name}-Card_{GPU_num}-Data_{data_now}")
    
    print(f"\033[36m训练指令：python tools/train.py {alg_file_pth} --work-dir {work_dir_file}")


    