# optimizer（优化器设计）TODO
# type_of_back_bone = "Vit"
def generate_optim_wrapper(type_of_back_bone=None):

    # optim_wrapper 配置-1
    optim_wrapper_1 = {
        'type': 'OptimWrapper',          # 表示这是一个优化器包装器（OptimWrapper）
        '_delete_': True,                # 通常用于在继承配置时删除旧的优化器配置，替换为新的优化器配置
        'optimizer': {
            'type': 'AdamW',             # 优化器类型为 AdamW
            'lr': 0.0001,                # 学习率，通常设定为一个较小的值
            'weight_decay': 0.0005       # 权重衰减系数
        },
        'clip_grad': {
            'max_norm': 1,               # 梯度裁剪的最大范数
            'norm_type': 2               # L2 范数
        }
    }
    # optim_wrapper 配置-2
    optim_wrapper_2 = {
        'type': 'OptimWrapper',          # 表示这是一个优化器包装器（OptimWrapper）
        '_delete_': True,                # 通常用于在继承配置时删除旧的优化器配置，替换为新的优化器配置
        'optimizer': {
            'type': 'SGD',               # 优化器类型为 SGD
            'lr': 0.05,                  # 学习率
            'weight_decay': 0.0005,       # 权重衰减系数
            'momentum': 0.9
        },
        'clip_grad': {
            'max_norm': 1,               # 梯度裁剪的最大范数
            'norm_type': 2               # L2 范数
        }
    }

    optim_wrapper_list = [optim_wrapper_1, optim_wrapper_2]

    # 打印所有可用的优化器选项
    while True:
        print("请选择 optim_wrapper (按 Enter 使用默认):")
        for i, scheduler in enumerate(optim_wrapper_list, 1):
            optimizer_type = scheduler['optimizer']['type']
            learning_rate = scheduler['optimizer']['lr']
            print(f"{i}. Optimizer: {optimizer_type}, LR: {learning_rate}")

            # 如果是vit网络，则需要额外的paramwise_cfg
            if type_of_back_bone != None and (type_of_back_bone.lower() == 'vit' or type_of_back_bone.lower() == 'visiontransformer'):
                custom_keys={
                    'pos_embed': dict(decay_mult=0.), # 位置嵌入（positional embeddings）。decay_mult=0. 意味着对这些嵌入不应用权重衰减
                    'cls_token': dict(decay_mult=0.), # 是在某些模型（如 Transformer 或 BERT）中，用于分类任务的特定 token
                    'norm': dict(decay_mult=0.) # 对归一化层的参数也禁用了权重衰减
                }
                optim_wrapper_list[i-1]['paramwise_cfg'] = dict(custom_keys)
                optim_wrapper_list[i-1]['_delete_'] = True
            
            if type_of_back_bone != None and (type_of_back_bone.lower() == 'swin'):
                custom_keys={
                    'absolute_pos_embed': dict(decay_mult=0.),
                    'relative_position_bias_table': dict(decay_mult=0.),
                    'norm': dict(decay_mult=0.)
                }
                optim_wrapper_list[i-1]['paramwise_cfg'] = dict(custom_keys)
                optim_wrapper_list[i-1]['_delete_'] = True

        choice = input(f"输入 1 到 {len(optim_wrapper_list)} 进行选择（默认 1）: ")

        # 如果用户按下 Enter 或选择 1，默认返回 optim_wrapper_1
        if choice == '' or choice == '1':
            return optim_wrapper_list[0]
        elif choice.isdigit() and 1 <= int(choice) <= len(optim_wrapper_list):
            return optim_wrapper_list[int(choice) - 1]
        else:
            print(f"无效输入，请输入 1 到 {len(optim_wrapper_list)} 或按 Enter 使用默认值")

def generate_param_scheduler(train_type, train_time_or_epoch):
    """
    根据训练模式（epoch或iteration）动态生成并选择学习率调度器。

    Args:
        train_type (str): 训练模式, 'epoch' 或 'iteration'。
        train_time_or_epoch (int): 训练的总轮数(epochs)或总迭代次数(k)。
    """
    # 1. 根据训练模式确定核心参数
    if train_type.lower() == 'epoch':
        by_epoch = True
        end_value = train_time_or_epoch
        warmup_end = 10  # 为epoch模式设置一个合理的10个epoch的预热期
        # 按比例调整MultiStepLR的milestones
        milestones = [int(end_value * 0.75), int(end_value * 0.9)]
    else:  # 默认为 iteration 模式
        by_epoch = False
        end_value = train_time_or_epoch * 1000
        warmup_end = 1500  # iteration模式使用原有的1500次迭代预热
        # MultiStepLR的milestones, 同样适配总时长
        milestones = [int(end_value * 0.75), int(end_value * 0.9)]
        
    # 2. 使用动态参数定义调度器模板
    # param_scheduler 配置-1
    param_scheduler_1 = [
        dict(
            type='LinearLR',
            start_factor=1e-6,
            by_epoch=by_epoch,
            begin=0,
            end=warmup_end),
        dict(
            type='PolyLR',
            power=0.9,
            begin=warmup_end,
            end=end_value,
            eta_min=1e-5,
            by_epoch=by_epoch)
    ]
    # param_scheduler 配置-2
    param_scheduler_2 = [
        dict(
            type='LinearLR',
            start_factor=1e-6,
            by_epoch=by_epoch,
            begin=0,
            end=warmup_end),
        dict(
            type='PolyLR',
            power=1.0,
            begin=warmup_end,
            end=end_value,
            eta_min=0.0,
            by_epoch=by_epoch)
    ]
    # param_scheduler 配置-3
    param_scheduler_3 = [
        dict(
            type='LinearLR', 
            start_factor=0.001, 
            by_epoch=by_epoch, 
            begin=0,
            end=warmup_end),
        dict(
            type='MultiStepLR',
            begin=warmup_end,
            end=end_value,
            milestones=milestones,
            by_epoch=by_epoch)
    ]
    param_scheduler_list = [param_scheduler_1, param_scheduler_2, param_scheduler_3]

    # 3. 打印并让用户选择 (这部分逻辑不变)
    while True:
        print("请选择 param_scheduler (学习率调度器):")
        for i, schedulers in enumerate(param_scheduler_list, 1):
            print(f"Scheduler {i}:")
            for scheduler in schedulers:
                # 提取字段并处理缺省情况
                scheduler_type = scheduler.get('type', '/')
                begin = scheduler.get('begin', '/')
                end = scheduler.get('end', '/')
                power = scheduler.get('power', '/')
                eta_min = scheduler.get('eta_min', '/')
                milestones_val = scheduler.get('milestones', '/')
                
                # 根据类型决定显示哪些信息
                if scheduler_type == 'MultiStepLR':
                    print(f"  - {scheduler_type}: begin={begin}, end={end}, milestones={milestones_val}")
                else:
                    print(f"  - {scheduler_type}: begin={begin}, end={end}, power={power}, eta_min={eta_min}")
        
        choice = input(f"输入 1 到 {len(param_scheduler_list)} 进行选择（默认 1）: ").strip()
        
        if choice == '' or choice == '1':
            return param_scheduler_list[0]
        elif choice.isdigit() and 1 <= int(choice) <= len(param_scheduler_list):
            return param_scheduler_list[int(choice) - 1]
        else:
            print(f"无效输入，请输入 1 到 {len(param_scheduler_list)} 或按 Enter 使用默认值\n")