# 选择crop_size大小
def select_crop_size(predefined_options=[(512, 512), (512, 1024), (769, 769)]):
    """
    让用户选择裁剪大小。默认选择 (512, 512)，用户可以选择其他预定义的裁剪大小，
    或者输入自定义的大小，且自定义的数字不能为负数。
    
    :return: 选择的裁剪大小 (width, height)
    """
    # 预定义的裁剪大小选项
    predefined_options = {str(i+1): option for i, option in enumerate(predefined_options)}
    
    # 显示可选的裁剪大小
    print("可用的裁剪大小选项：")
    for key, value in predefined_options.items():
        print(f"{key}. {value}", end="  ")
    print(f"{len(predefined_options)+1}. 自定义大小")
    
    # 用户选择
    choice = input("请选择裁剪大小选项 (默认 1): ").strip()
    
    # 如果用户没有输入，使用默认值
    if choice == "" or choice == "1":
        return predefined_options["1"]
    
    # 如果用户选择了预定义的选项
    if choice in predefined_options:
        return predefined_options[choice]
    
    # 如果用户选择自定义大小
    if choice == f"{len(predefined_options)+1}":
        while True:
            try:
                width = int(input("请输入裁剪宽度 (正整数): "))
                height = int(input("请输入裁剪高度 (正整数): "))
                if width > 0 and height > 0:
                    return (width, height)
                else:
                    print("宽度和高度必须是正整数。")
            except ValueError:
                print("输入无效，请输入有效的正整数。")
    
    # 如果用户输入无效，返回默认值
    print("无效选择，使用默认值 (512, 512)")
    return predefined_options["1"]

# 大字典，包含所有模型的信息
pretrained_models_dict = {
    "openmmlab/resnet18_v1c": {'pth': './My_Local_Model/open_mmlab/resnet18_v1c.pth', 'depth': 18, 'type':'ResNetV1c'},
    "openmmlab/resnet50_v1c": {'pth': './My_Local_Model/open_mmlab/resnet50_v1c.pth', 'depth': 50, 'type':'ResNetV1c'},
    "openmmlab/resnet101_v1c": {'pth': './My_Local_Model/open_mmlab/resnet101_v1c.pth', 'depth': 101, 'type':'ResNetV1c'},
    "torchvision://resnet18": {'pth': './My_Local_Model/torchvision_012/resnet18.pth', 'depth': 18, 'type':'ResNet'},
    "torchvision://resnet50": {'pth': './My_Local_Model/torchvision_012/resnet50.pth', 'depth': 50, 'type':'ResNet'},
    "torchvision://resnet101": {'pth': './My_Local_Model/torchvision_012/resnet101.pth', 'depth': 101, 'type':'ResNet'},
    "openmmlab/pidnet-s":{'pth': './My_Local_Model/open_mmlab/pidnet-s.pth', 'type':'pidnet', 'size':'small'},
    "openmmlab/pidnet-m":{'pth': './My_Local_Model/open_mmlab/pidnet-m.pth', 'type':'pidnet', 'size':'medium'},
    "openmmlab/pidnet-l":{'pth': './My_Local_Model/open_mmlab/pidnet-l.pth', 'type':'pidnet', 'size':'large'},
    "openmmlab/ddrnet23-s":{'pth': './My_Local_Model/open_mmlab/ddrnet23-s.pth', 'type':'ddrnet', 'size':'small'},
    "openmmlab/ddrnet23":{'pth': './My_Local_Model/open_mmlab/ddrnet23.pth', 'type':'ddrnet', 'size':'normal'},
    "openmmlab/stdc1":{'pth': './My_Local_Model/open_mmlab/stdc1.pth', 'type':'stdc', 'size':'V1'},
    "openmmlab/stdc2":{'pth': './My_Local_Model/open_mmlab/stdc2.pth', 'type':'stdc', 'size':'V2'},
    "pretrain/vit-b16_p16_224-80ecf9dd.pth":{'pth': './My_Local_Model/pretrain/vit-b16_p16_224-80ecf9dd.pth'},
    "pretrain/beit_base_patch16_224_pt22k_ft22k.pth":{'pth': './My_Local_Model/pretrain/beit_base_patch16_224_pt22k_ft22k.pth'},
    "pretrain/beit_large_patch16_224_pt22k_ft22k.pth":{'pth': './My_Local_Model/pretrain/beit_large_patch16_224_pt22k_ft22k.pth'},
    "pretrain/swin_large-d5bdebaf.pth":{'pth': './My_Local_Model/pretrain/swin_large_patch4_window7_224_22k_20220308-d5bdebaf.pth', 'type':'swin', 'size':'large'},
    "pretrain/swin_tiny-f41b89d3.pth":{'pth': './My_Local_Model/pretrain/swin_tiny_patch4_window7_224_20220308-f41b89d3.pth', 'type':'swin', 'size':'tiny'},
    "mae_pretrain_vit_base_mmcls.pth":{'pth': './My_Local_Model/pretrain/mae_pretrain_vit_base_mmcls.pth'},
    "open-mmlab://msra/hrnetv2_w18":{'pth': './My_Local_Model/open_mmlab/msra/hrnetv2_w18.pth'},
    "open-mmlab://msra/hrnetv2_w18_small":{'pth': './My_Local_Model/open_mmlab/msra/hrnetv2_w18_small.pth'},
    'open-mmlab://msra/hrnetv2_w48':{'pth': './My_Local_Model/open_mmlab/msra/hrnetv2_w48.pth'},
    "pretrain/swin_large-6580f57d.pth":{'pth': './My_Local_Model/pretrain/swin_large_patch4_window12_384_22k_20220412-6580f57d.pth', 'type':'swin', 'size':'large'}, # mask2former
    "pretrain/swin_base-e5c09f74.pth":{'pth': './My_Local_Model/pretrain/swin_base_patch4_window12_384_22k_20220317-e5c09f74.pth', 'type':'swin', 'size':'base'}, # mask2former
    "pretrain/swin_small-7ba6d6dd.pth":{'pth': './My_Local_Model/pretrain/swin_small_patch4_window7_224_20220317-7ba6d6dd.pth', 'type':'swin', 'size':'small'}, # mask2former
    "pretrain/swin_tiny-1cdeb081.pth":{'pth': './My_Local_Model/pretrain/swin_tiny_patch4_window7_224_20220317-1cdeb081.pth', 'type':'swin', 'size':'tiny'}, # mask2former
    # 这里可以包含更多模型信息...
}

# 选择预训练模型
def select_pretrained_model(model_list, need_select_pretrained = False, pretrained_models_dict = pretrained_models_dict):
    """
    让用户从给定的模型列表中选择预训练模型、是否选择预训练（否-默认开启预训练），并返回对应的 pth 路径和 其他 信息。
    
    :param model_list: 可用的模型名称列表，例如 ['openmmlab/resnet50_v1c', 'openmmlab/resnet101_v1c']
    :return: (pretrained_pth, depth)
    """
    
    # 过滤传入的模型列表，确保它们在字典中有信息
    valid_models = {key: pretrained_models_dict[key] for key in model_list if key in pretrained_models_dict}

    if not valid_models:
        print("错误：提供的模型列表中没有有效的模型信息。")
        return None, None
    
    # 显示可用的预训练模型
    print("可用的预训练模型类型：")
    for i, (model_name, model_info) in enumerate(valid_models.items(), 1):
        print(f"{i}. {model_name}")

    # 用户选择
    choice = input(f"请选择预训练模型编号 (1-{len(valid_models)}, 默认 1): ").strip()

    # 如果用户没有输入，或输入无效，使用默认值
    if not choice.isdigit() or not (1 <= int(choice) <= len(valid_models)):
        choice = "1"
    
    # 获取用户选择的模型信息
    selected_model_name = list(valid_models.keys())[int(choice) - 1]
    selected_model_info = valid_models[selected_model_name]

    # TODO 选择特定信息
    # # 返回模型的 pth 路径和 depth
    pretrained_pth = selected_model_info['pth']
    # depth = selected_model_info['depth']
    
    # print(f"    已选择模型: {selected_model_name} depth: {depth} pth: {pretrained_pth}")
    print(f"    已选择模型: {selected_model_name} pth: {pretrained_pth}")

    if need_select_pretrained == False:
        print("默认开启预训练")
        select_pretrained = True
        return selected_model_name, select_pretrained, pretrained_pth, selected_model_info
    else:
        # 提示用户输入 True 或 False，或者直接按 Enter 默认使用预训练模型    
        choice = input("是否使用预训练模型？输入 Y（使用）或 N（不使用），直接按 Enter 默认使用：True：")
        while True:
            # 如果用户没有输入，默认使用预训练模型
            if choice == '':
                select_pretrained = True
                break
            # 转换输入为布尔值
            elif choice.lower() == 'y':
                select_pretrained = True
                break
            elif choice.lower() == 'n':
                select_pretrained = False
                break
            else:
                print("无效输入，请输入 'True' 或 'False'，或直接按 Enter 选择默认值")
            
        return selected_model_name, select_pretrained, pretrained_pth, selected_model_info

# 大字典，包含所有模型的信息
samplers_dict = {
    "OHEMPixelSampler": dict(type='OHEMPixelSampler', thresh=0.7, min_kept=10000),
    # 这里可以包含更多采样函数模型信息...
}

# 选择采样函数
def select_sampler(sampler_list, samplers_dict = samplers_dict):

    # 提示用户输入 True 或 False，或者直接按 Enter 默认使用预训练模型    
    choice = input("是否使用采样函数？输入 Y（使用）或 N（不使用），直接按 Enter 默认不使用：False：")
    while True:
        # 如果用户没有输入，默认使用预训练模型
        if choice == '':
            use_sampler = False
            return None, use_sampler, None
        # 转换输入为布尔值
        elif choice.lower() == 'y':
            use_sampler = True
            break
        elif choice.lower() == 'n':
            use_sampler = False
            return None, use_sampler, None
        else:
            print("无效输入，请输入 'True' 或 'False'，或直接按 Enter 选择默认值")

    # 过滤传入的模型列表，确保它们在字典中有信息
    valid_samplers = {key: samplers_dict[key] for key in sampler_list if key in samplers_dict}

    if not valid_samplers:
        print("错误：提供的采样函数列表中没有有效的采样函数信息。")
        return None, None

    # 如果只有一个可用的采样函数，直接选择
    if len(valid_samplers) != 1:
        # 显示可用的采样函数
        print("可用的采样函数：")
        for i, (sampler_name, sampler_info) in enumerate(valid_samplers.items(), 1):
            print(f"{i}. {sampler_name}")

        # 用户选择
        choice = input(f"请选择采样函数编号 (1-{len(valid_samplers)}, 默认 1): ").strip()

        # 如果用户没有输入，或输入无效，使用默认值
        if not choice.isdigit() or not (1 <= int(choice) <= len(valid_samplers)):
            choice = "1"
    else:
        choice = "1"
    
    # 获取用户选择的模型信息
    selected_sampler_name = list(valid_samplers.keys())[int(choice) - 1]
    selected_sampler_info = valid_samplers[selected_sampler_name]

    print(f"    已选择采样函数: {selected_sampler_name}")

    return selected_sampler_name, use_sampler, selected_sampler_info

# 选择test_cfg中是否滑动，是否默认选择select_slide
def select_test_cfg_slide(crop_size, select_slide=False):
    """
    让用户选择是否使用滑动窗口，并根据选择设置相应的模式和参数。
    
    :param crop_size: 输入的裁剪大小 (宽, 高)
    :return: test_cfg_mode, test_cfg_crop_div_stride, crop_size
    """
    if select_slide == False:
        # 提示用户是否选择滑动窗口模式
        use_slide = input("是否选择滑动窗口模式？(y/n, 默认 n): ").strip().lower()
    else:
        use_slide = 'y' # 使用滑动窗口模式

    # 默认不使用滑动窗口模式
    if use_slide == 'y':
        test_cfg_mode = 'slide'
        # 提示用户输入 test_cfg_crop_div_stride，默认值为 1.5
        try:
            test_cfg_crop_div_stride = input("请输入滑动窗口的 stride 除以 crop_size 比例 (默认 1.5): ").strip()
            test_cfg_crop_div_stride = float(test_cfg_crop_div_stride) if test_cfg_crop_div_stride else 1.5
        except ValueError:
            print("输入无效，使用默认比例 1.5")
            test_cfg_crop_div_stride = 1.5

        # 计算 stride
        stride = tuple(int(c / test_cfg_crop_div_stride) for c in crop_size)
        print(f"  已选择滑动窗口模式: {test_cfg_mode}")
        print(f"crop_size: {crop_size}")
        print(f"stride: {stride}")
    else:
        test_cfg_mode = None  # 默认模式
        test_cfg_crop_div_stride = None
        stride = None
        print(f"  关闭滑动窗口模式")
    
    return test_cfg_mode, test_cfg_crop_div_stride


