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, select_test_cfg_slide

from Initial_Alg_Program.Initial_Alg_Gen_Tool import format_all_data, write_config_to_file, create_dict_by_kwargs, 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

if __name__ == '__main__':
    ########### 1.超参数-算法默认参数 ###########
    alg_file_name = 'knet_r50-d8_my'

    ########### 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. 设定参数 ########### 

    _base_ = generate_base_config(alg_file_name=alg_file_name, dataset_file_name=dataset_file_name, schedule_file_name=schedule_file_name) # _base_无算法

    crop_size = select_crop_size(predefined_options=[(512, 512), (640,640)]) # 选择切割大小

    data_preprocessor = generate_data_preprocessor(crop_size = crop_size, mean = mean, std = std)
    
    model_data_preprocessor = data_preprocessor # 修改data_preprocessor
    model_data_preprocessor = generate_model_data_preprocessor(model_data_preprocessor=model_data_preprocessor)
    
    norm_cfg = generate_norm_cfg(GPU_num=GPU_num)
    
    # 选择对应的模型
    config_dict = {
        'DeepLabV3': {
            'decode_head_kernel_generate_head_dict': dict(
                _delete_ = True,
                type='ASPPHead',
                in_channels=2048,
                in_index=3,
                channels=512,
                dilations=(1, 12, 24, 36),
                dropout_ratio=0.1,
                num_classes=150,
                norm_cfg=norm_cfg,
                align_corners=False),
            'backbone_dilations': (1, 1, 2, 4),
            'backbone_strides': (1, 2, 1, 1),
            'auxiliary_head_in_channels':1024
        },
        'PSPNet': {
            'decode_head_kernel_generate_head_dict': dict(
                _delete_ = True,
                type='PSPHead',
                in_channels=2048,
                in_index=3,
                channels=512,
                pool_scales=(1, 2, 3, 6),
                dropout_ratio=0.1,
                num_classes=150,
                norm_cfg=norm_cfg,
                align_corners=False),
            'backbone_dilations': (1, 1, 2, 4),
            'backbone_strides': (1, 2, 1, 1),
            'auxiliary_head_in_channels':1024
        },
        'FCN': {
            'decode_head_kernel_generate_head_dict': dict(
                _delete_ = True,
                type='FCNHead',
                in_channels=2048,
                in_index=3,
                channels=512,
                num_convs=2,
                concat_input=True,
                dropout_ratio=0.1,
                num_classes=150,
                norm_cfg=norm_cfg,
                align_corners=False),
            'backbone_dilations': (1, 1, 2, 4),
            'backbone_strides': (1, 2, 1, 1),
            'auxiliary_head_in_channels':1024
        },
        'UPerNet': {
            'decode_head_kernel_generate_head_dict': dict(
                type='UPerHead',
                in_channels=[256, 512, 1024, 2048],
                in_index=[0, 1, 2, 3],
                pool_scales=(1, 2, 3, 6),
                channels=512,
                dropout_ratio=0.1,
                num_classes=150,
                norm_cfg=norm_cfg,
                align_corners=False),
            'backbone_dilations': (1, 1, 1, 1),
            'backbone_strides': (1, 2, 2, 2),
            'auxiliary_head_in_channels':1024
        }
    }

    models = ['DeepLabV3', 'PSPNet', 'FCN', 'UPerNet']
    print("请你选择对应的模型:")
    while True:
        for index, model in enumerate(models, 1):
            print(f"{index}. {model}")
        try:
            user_input = int(input("Enter your choice (1-4): "))
            if 1 <= user_input <= 4:
                model_choice = models[user_input - 1]
                break
            else:
                print("Invalid input, please enter a number between 1 and 4.")
        except ValueError:
            print("Invalid input, please enter a valid number.")

    decode_head_kernel_generate_head_dict = config_dict[model_choice]['decode_head_kernel_generate_head_dict']
    backbone_dilations = config_dict[model_choice]['backbone_dilations']
    backbone_strides = config_dict[model_choice]['backbone_strides']
    auxiliary_head_in_channels = config_dict[model_choice]['auxiliary_head_in_channels']

    model_list = ['openmmlab/resnet50_v1c', 'openmmlab/resnet101_v1c', "pretrain/swin_tiny-f41b89d3.pth", "pretrain/swin_large-d5bdebaf.pth"] 
    selected_model_name, select_pretrained, pretrained_pth, selected_model_info = select_pretrained_model(model_list=model_list)
    if selected_model_info['type'] == 'ResNetV1c':
        depth = selected_model_info['depth']
        backbone = create_dict_by_kwargs(depth = depth, norm_cfg=norm_cfg, dilations=backbone_dilations, strides=backbone_strides)
    elif  selected_model_info['type'] == 'swin':
        if selected_model_name == 'pretrain/swin_tiny-f41b89d3.pth':
            backbone = dict(
                _delete_=True,
                type='SwinTransformer',
                embed_dims=96, # 
                depths=[2, 2, 6, 2], # 
                num_heads=[3, 6, 12, 24], #
                window_size=7,
                mlp_ratio=4,
                qkv_bias=True,
                qk_scale=None,
                drop_rate=0.,
                attn_drop_rate=0.,
                drop_path_rate=0.3, # 
                use_abs_pos_embed=False,
                patch_norm=True,
                out_indices=(0, 1, 2, 3))
        if selected_model_name == 'pretrain/swin_large-d5bdebaf.pth':
            backbone = dict(
                _delete_=True,
                type='SwinTransformer',
                embed_dims=192, # 
                depths=[2, 2, 18, 2], # 
                num_heads=[6, 12, 24, 48], # 
                window_size=7,
                mlp_ratio=4,
                qkv_bias=True,
                qk_scale=None,
                drop_rate=0.,
                attn_drop_rate=0.,
                drop_path_rate=0.4, # 
                use_abs_pos_embed=False,
                patch_norm=True,
                out_indices=(0, 1, 2, 3))
        backbone.update(create_dict_by_kwargs(norm_cfg=dict(type='LN'))) # TODO Transformer一类的内容norm_cfg都设置为LN TODO  # 没有 dilations 和 strides
        
    test_cfg_mode, test_cfg_crop_div_stride = select_test_cfg_slide(crop_size)

    # 3.2. 自动生成参数
    # 如果使用slide模式，则开启align_corners
    if test_cfg_mode == "slide":
        align_corners = True
        test_slide = 'testslide' # 算法名进行标记
    else:
        align_corners = False
        test_slide = 'no_testslide' # 算法名进行标记


    ########### 3.生成各个部分 ########### 

    # 3.2. decode_head
    decode_head = get_var_from_py_file(os.path.join('./configs/_base_/models', alg_file_name+'.py'), 'model')['decode_head']

    decode_head_kernel_update_head_list = decode_head['kernel_update_head']
    for i in range(len(decode_head_kernel_update_head_list)):
        decode_head_kernel_update_head_list[i]['num_classes'] = num_classes
    decode_head['kernel_update_head'] = decode_head_kernel_update_head_list

    decode_head_kernel_generate_head_dict['num_classes'] = num_classes
    if  model_choice == 'UPerNet':
        if selected_model_name == 'pretrain/swin_tiny-f41b89d3.pth':
            decode_head_kernel_generate_head_dict['in_channels'] = [96, 192, 384, 768]
            auxiliary_head_in_channels = 384
        elif selected_model_name == 'pretrain/swin_large-d5bdebaf.pth':
            decode_head_kernel_generate_head_dict['in_channels'] = [192, 384, 768, 1536]
            auxiliary_head_in_channels = 768
    elif model_choice in ['DeepLabV3', 'PSPNet', 'FCN']:
        if selected_model_name == 'pretrain/swin_tiny-f41b89d3.pth':
            decode_head_kernel_generate_head_dict['in_channels'] = 768
            auxiliary_head_in_channels = 384
        elif selected_model_name == 'pretrain/swin_large-d5bdebaf.pth':
            decode_head_kernel_generate_head_dict['in_channels'] = 1536
            auxiliary_head_in_channels = 768
        

    decode_head_kernel_generate_head_dict_loss_decode = dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0) # Way Ori
    decode_head_kernel_generate_head_dict_loss_decode = dict(type='DiceLoss', use_sigmoid=False, loss_weight=1.0) # Way 1
    decode_head_kernel_generate_head_dict['loss_decode'] = decode_head_kernel_generate_head_dict_loss_decode

    decode_head['kernel_generate_head'] = decode_head_kernel_generate_head_dict

    decode_head['kernel_generate_head']['align_corners'] = align_corners

    # 3.2. auxiliary_head
    auxiliary_head = get_var_from_py_file(os.path.join('./configs/_base_/models', alg_file_name+'.py'), 'model')['auxiliary_head']
    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
    auxiliary_head['loss_decode'] = auxiliary_head_loss_decode
    auxiliary_head['align_corners'] = align_corners
    auxiliary_head['in_channels'] = auxiliary_head_in_channels

    # 综合model
    model = dict(
        backbone = backbone,
        pretrained = pretrained_pth,
        data_preprocessor = model_data_preprocessor,
        decode_head = decode_head,
        auxiliary_head = auxiliary_head,
    )

    # 4. train_dataloader
    train_dataloader, batch_size, num_workers = generate_train_dataloader(batch_size_default=2, num_workers_default=2)
    if selected_model_info['type'] == 'ResNetV1c':
        optim_wrapper = generate_optim_wrapper()
    elif selected_model_info['type'] == 'swin':
        optim_wrapper = generate_optim_wrapper('swin')
        
    param_scheduler = generate_param_scheduler(train_type, train_time_or_epoch_k)


    ########### 4.程序写入 ###########
    # 算法名称解析：knet【Alg】 _ r50【pretrained模型深度】 - d8【_base_/models】 _ 4【GPU数量】 x b2【Batch Size大小】 - 40k【schedule】 _ cityscapes【数据集】 - 512x1024【crop_size】.py
    if selected_model_info['type'] == 'ResNetV1c':
        alg_file_name = f"{alg_name}_r{depth}_{model_choice}_{'Pre' if select_pretrained else 'NoPre'}_g{GPU_num}-{schedule_file_name.lstrip('schedule_')}_{dataset_file_name}-{crop_size[0]}x{crop_size[1]}-{test_slide}.py"
    elif selected_model_info['type'] == 'swin':
        alg_file_name = f"{alg_name}_swin_{selected_model_info['size']}_{model_choice}_{'Pre' if select_pretrained else 'NoPre'}_g{GPU_num}-{schedule_file_name.lstrip('schedule_')}_{dataset_file_name}-{crop_size[0]}x{crop_size[1]}-{test_slide}.py"
    output_configs_alg_my_alg = os.path.join(f'./configs/knet/', 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, train_dataloader=train_dataloader, optim_wrapper=optim_wrapper, param_scheduler=param_scheduler)

    # 将信息写入文件
    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)
