#!/usr/bin/env python3
# -*- coding: UTF-8 -*-
import os
from Initial_Alg_Program.Initial_Alg_Gen_Tool import format_all_data, write_config_to_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

if __name__ == '__main__':
    ########### 1.1.定义算法基本参数 ###########
    # A. generate_base_config
    alg_name = 'ann' # ./configs中算法简称
    alg_file_name = 'ann_r50-d8' # 算法根文件
    dataset_file_name = 'my_dataset_model' # 数据文件
    schedule_file_name = 'schedule_4k_check_400' # schedule文件
    _base_ = generate_base_config(alg_file_name=alg_file_name, dataset_file_name=dataset_file_name, schedule_file_name=schedule_file_name)

    # B. generate_norm_cfg
    GPU_num = 1 # GPU数量 
    norm_cfg = generate_norm_cfg(GPU_num=GPU_num)

    # C. generate_data_preprocessor
    crop_size = (512,512) # 分割大小
    data_preprocessor = generate_data_preprocessor(crop_size = crop_size)

    # D. generate_model
        # D.1. pretrained
    pretrained_pth = './My_Local_Model/open_mmlab/resnet50_v1c.pth' # 预训练模型位置
    pretrained = generate_model_pretrained(pretrained_pth = pretrained_pth)
    
        # D.2. backbone
    depth = 50 # 模型深度
    backbone = generate_model_backbone(depth=depth)

        # D.3. data_preprocessor
    model_data_preprocessor = data_preprocessor # 修改data_preprocessor
    model_data_preprocessor = generate_model_data_preprocessor(model_data_preprocessor=model_data_preprocessor)

        # D.4. decode_head、auxiliary_head
    num_classes=36 # 分类数
    decode_head_loss_decode_dict=dict(type='DiceLoss', use_sigmoid=False, loss_weight=1.0) # decode损失函数
    align_corners=False # 是否需要角对齐
    auxiliary_head_loss_decode_dict=dict(type='DiceLoss', use_sigmoid=False, loss_weight=0.4) # DiceLoss损失函数
    decode_head = generate_model_decode_head(num_classes=num_classes, decode_head_loss_decode_dict=decode_head_loss_decode_dict, align_corners=align_corners)
    auxiliary_head = generate_model_auxiliary_head(num_classes=num_classes, auxiliary_head_loss_decode_dict=auxiliary_head_loss_decode_dict, align_corners=align_corners)

        # D.5. train_cfg
    # train_cfg = generate_model_train_cfg()

        # D.6. test_cfg
    # test_cfg = generate_model_test_cfg()
    
    # E. 综合model
    model = dict(
        pretrained = pretrained,
        backbone = backbone,
        data_preprocessor = model_data_preprocessor,
        decode_head = decode_head,
        auxiliary_head = auxiliary_head,
        # train_cfg =  train_cfg,
        # test_cfg = test_cfg,
    )

    ########### 2.文件存储 ###########
    # output_configs_alg_my_alg = os.path.join(f'my_{alg_name}.py')
    output_configs_alg_my_alg = os.path.join(f'./configs/{alg_name}/', f'my_{alg_name}.py')
    
    write_config_to_file(output_configs_alg_my_alg, _base_=_base_, crop_size=crop_size, data_preprocessor=data_preprocessor, model=model)

