import os, sys, argparse, json
# 添加上级目录到 sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../')))

from Initial_Alg_Program.Initial_Alg_Select_Tool import select_crop_size, select_pretrained_model

from Initial_Alg_Program.Initial_Alg_Gen_Tool import format_all_data, write_config_to_file, create_dict_by_kwargs, get_var_from_file, update_list_dict_var, get_var_from_py_file
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
from Initial_Alg_Program.Initial_Alg_Gen_4_optimizer import generate_optim_wrapper, generate_param_scheduler
from Initial_Alg_Program.Initial_Alg_Gen_5_train_dataloader import generate_train_dataloader

# 交互式选择 decode_head 的函数
def select_decode_head(decode_head_choose):
    print("可用的 decode head 选项：")
    for i, key in enumerate(decode_head_choose.keys()):
        print(f"{i + 1}. {key}")

    while True:
        try:
            # 提示用户输入选择
            choice = int(input("请选择需要的 decode head（输入编号）："))
            # 检查输入是否在有效范围内
            if 1 <= choice <= len(decode_head_choose):
                selected_key = list(decode_head_choose.keys())[choice - 1]
                print(f"你选择了: {selected_key}")
                return decode_head_choose[selected_key]
            else:
                print("输入的编号不正确，请重新输入。")
        except ValueError:
            print("输入无效，请输入有效的编号。")

if __name__ == '__main__':
    ########### 1.超参数-算法默认参数 ###########
    alg_file_name = 'upernet_beit'
    alg_file_pth = f"configs/_base_/models/{alg_file_name}.py"

    ########### 2.参数解析 ########### 
    parser = argparse.ArgumentParser(description="Run algorithm with specified parameters.")

    # 添加命令行参数
    parser.add_argument('--alg_name', type=str, required=True, help="Algorithm name (e.g., 'ann')")
    parser.add_argument('--dataset_file_name', type=str, required=True, help="Dataset file name (e.g., 'my_dataset_model')")
    parser.add_argument('--mean', type=str, required=True, help="Mean of dataset (e.g., [1.2, 1.2, 1.2])")
    parser.add_argument('--std', type=str, required=True, help="Mean of dataset (e.g., [1.2, 1.2, 1.2])")
    parser.add_argument('--dataset_num_classes', type=int, required=True, help="Dataset class num (e.g., 36)")
    parser.add_argument('--GPU_num', type=int, required=True, help="Num of GPU to train (e.g., 1)")
    parser.add_argument('--schedule_file_name', type=str, required=True, help="Schedule file name (e.g., 'schedule_4k_check_400')")
    parser.add_argument('--train_type', type=str, required=True, help="Schedule train type (e.g., Epoch or Iteration)")
    parser.add_argument('--train_time_or_epoch_k', type=int, required=True, help="Schedule train time k (e.g., 4)")

    # 解析命令行参数
    args = parser.parse_args()

    # 变量赋值
    alg_name = args.alg_name
    dataset_file_name = args.dataset_file_name
    mean = json.loads(args.mean) # 使用json.loads转换成list
    std = json.loads(args.std) # 使用json.loads转换成list
    num_classes = args.dataset_num_classes
    schedule_file_name = args.schedule_file_name
    train_type = args.train_type
    train_time_or_epoch_k = args.train_time_or_epoch_k
    GPU_num = args.GPU_num
    
    ########### 3.1. 生成_base_部分 ########### 
    _base_ = generate_base_config(alg_file_name=alg_file_name, dataset_file_name=dataset_file_name, schedule_file_name=schedule_file_name)
    
    ########### 3.1. 生成norm_cfg部分 ########### 
    norm_cfg = generate_norm_cfg(GPU_num=GPU_num)

    ########### 3.2. 生成data_preprocessor部分 ########### 
    crop_size = select_crop_size(predefined_options=[(640, 640)]) # 选择切割大小
    data_preprocessor = generate_data_preprocessor(crop_size = crop_size, mean = mean, std = std)

    ########### 3.3. 生成model部分 ########### 
    model_data_preprocessor = data_preprocessor # 修改data_preprocessor
    # 获取backbone模型、是否需要预训练
    model_list = ['pretrain/beit_base_patch16_224_pt22k_ft22k.pth', 'pretrain/beit_large_patch16_224_pt22k_ft22k.pth'] 
    selected_model_name, select_pretrained, pretrained_pth, selected_model_info = select_pretrained_model(model_list=model_list) # 需要选择是否用预训练模型

    backbone = create_dict_by_kwargs(type='BEiT', img_size=crop_size,)

    decode_head_loss_decode=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0) # Way Ori
    decode_head_loss_decode=dict(type='DiceLoss', use_sigmoid=False, loss_weight=1.0) # Way 1
    
    auxiliary_head_loss_decode=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4) # Way Ori
    auxiliary_head_loss_decode=dict(type='DiceLoss', use_sigmoid=False, loss_weight=0.4) # Way 1

    if selected_model_name == 'pretrain/beit_base_patch16_224_pt22k_ft22k.pth':
        backbone_new=dict(
            # patch_size=16,
            # in_channels=3,
            embed_dims=768,
            num_layers=12,
            num_heads=12,
            # mlp_ratio=4,
            out_indices=(3, 5, 7, 11),
            # qv_bias=True,
            # attn_drop_rate=0.0,
            drop_path_rate=0.1,
            # norm_cfg=dict(type='LN', eps=1e-6),
            # act_cfg=dict(type='GELU'),
            # norm_eval=False,
            init_values=0.1)
        neck = dict(type='Feature2Pyramid', embed_dim=768, rescales=[4, 2, 1, 0.5])
        decode_head=dict(in_channels=[768, 768, 768, 768], num_classes=num_classes, channels=768, norm_cfg=norm_cfg, loss_decode=decode_head_loss_decode)
        auxiliary_head=dict(norm_cfg=norm_cfg, in_channels=768, num_classes=num_classes, loss_decode = auxiliary_head_loss_decode)

        optim_wrapper = dict(
            _delete_=True,
            type='AmpOptimWrapper', # type='OptimWrapper', # TODO Ori
            optimizer=dict(
                type='AdamW', lr=3e-5, betas=(0.9, 0.999), weight_decay=0.05),
            constructor='LayerDecayOptimizerConstructor',
            paramwise_cfg=dict(num_layers=12, layer_decay_rate=0.9))
        
        model_size = 'base'

        train_dataloader, batch_size = generate_train_dataloader(2)

    elif selected_model_name == 'pretrain/beit_large_patch16_224_pt22k_ft22k.pth':
        backbone_new=dict(
            embed_dims=1024,
            num_layers=24,
            num_heads=16,
            # mlp_ratio=4,
            # qv_bias=True,
            init_values=1e-6,
            drop_path_rate=0.2,
            out_indices=[7, 11, 15, 23])
        neck=dict(type='Feature2Pyramid', embed_dim=1024, rescales=[4, 2, 1, 0.5])
        decode_head=dict(in_channels=[1024, 1024, 1024, 1024], num_classes=num_classes, channels=1024, norm_cfg=norm_cfg, loss_decode=decode_head_loss_decode)
        auxiliary_head=dict(norm_cfg=norm_cfg, in_channels=1024, num_classes=num_classes, loss_decode = auxiliary_head_loss_decode)

        optim_wrapper = dict(
            _delete_=True,
            type='AmpOptimWrapper',
            optimizer=dict(
                type='AdamW', lr=2e-5, betas=(0.9, 0.999), weight_decay=0.05),
            constructor='LayerDecayOptimizerConstructor',
            paramwise_cfg=dict(num_layers=24, layer_decay_rate=0.95),
            accumulative_counts=2)

        model_size = 'large'

        train_dataloader, batch_size = generate_train_dataloader(1)

    # 更新backbone
    backbone.update(backbone_new)

    model = dict(
        backbone = backbone,
        data_preprocessor = model_data_preprocessor,
        pretrained = pretrained_pth,
        neck = neck,
        decode_head = decode_head,
        auxiliary_head = auxiliary_head,
    )
    

    ########### 3.4. 生成optim、param_scheduler部分 ########### 
    optim_wrapper = optim_wrapper
    param_scheduler = generate_param_scheduler(train_type, train_time_or_epoch_k)

    ########### 4.程序写入 ###########
    # 算法名称解析：beit【Alg】 _ r50【pretrained模型深度】 - d8【_base_/models】 _ 4【GPU数量】 x b2【Batch Size大小】 - 40k【schedule】 _ cityscapes【数据集】 - 512x1024【crop_size】.py
    alg_file_name = f"{alg_name}-{model_size}_b{batch_size}_g{GPU_num}-{schedule_file_name.lstrip('schedule_')}_{dataset_file_name}-{crop_size[0]}x{crop_size[1]}.py"
    output_configs_alg_my_alg = os.path.join(f'./configs/beit/', alg_file_name)
    write_config_to_file(output_configs_alg_my_alg, _base_=_base_, norm_cfg=norm_cfg, crop_size=crop_size, data_preprocessor=data_preprocessor, model=model, optim_wrapper=optim_wrapper, param_scheduler=param_scheduler, train_dataloader = train_dataloader)

    # 将信息写入文件
    with open("_temp_.txt", "w") as file:
        json.dump({'alg_infos':{'alg_file_name':alg_file_name,'alg_file_pth':output_configs_alg_my_alg}}, file)
