from .Initial_Alg_Gen_Tool import get_var_from_file, format_dict

# 1. 修改pretrained
# pretrained_pth = './My_Local_Model/open_mmlab/resnet50_v1c.pth')
def generate_model_pretrained(pretrained_pth=None):
    pretrained = pretrained_pth
    if pretrained_pth == None:
        return None
    return pretrained

# 2. 修改backbone
# depth=50
def generate_model_backbone(depth=None):
    backbone = {}
    if depth != None:
        backbone['depth']=depth
    return backbone

# 3. 修改model_data_preprocessor
# model_data_preprocessor='data_preprocessor'
def generate_model_data_preprocessor(model_data_preprocessor=None):
    if model_data_preprocessor != None:
        model_data_preprocessor = model_data_preprocessor
    return model_data_preprocessor

# 4. 修改decode_head
# num_classes='36' # 分割数
# decode_head_loss_decode_dict=dict(type='DiceLoss', use_sigmoid=False, loss_weight=1.0) # 直接传入dict
# align_corners=False是否需要角对齐
def generate_model_decode_head(num_classes=None, decode_head_loss_decode_dict=None, align_corners=None):
    decode_head = {}
    if num_classes != None:
        decode_head['num_classes']=num_classes
    if decode_head_loss_decode_dict != None:
        decode_head['loss_decode']=decode_head_loss_decode_dict
    if align_corners != None:
        decode_head['align_corners']=align_corners
    
    return decode_head

# 5. 修改auxiliary_head
# num_classes='36' # 分割数
# auxiliary_head_loss_decode_dict=dict(type='DiceLoss', use_sigmoid=False, loss_weight=0.4) # 直接传入dict
# align_corners=False是否需要角对齐
def generate_model_auxiliary_head(num_classes=None, auxiliary_head_loss_decode_dict=None, align_corners=None):
    auxiliary_head = {}
    if num_classes != None:
        auxiliary_head['num_classes']=num_classes
    if auxiliary_head_loss_decode_dict != None:
        auxiliary_head['loss_decode']=auxiliary_head_loss_decode_dict
    if align_corners != None:
        auxiliary_head['align_corners']=align_corners
    
    return auxiliary_head

# 6. 修改train_cfg
def generate_model_train_cfg():
    train_cfg = {}
    
    return train_cfg

# 6. 修改test_cfg
# mode='slide'
# crop_size = (767,767)
# test_cfg_crop_div_stride = 1.5
def generate_model_test_cfg(test_cfg_mode=None, crop_size=None, test_cfg_crop_div_stride=None):
    test_cfg = {}
    if test_cfg_mode != None:
        test_cfg['mode'] = test_cfg_mode
    if crop_size != None:
        test_cfg['crop_size'] = crop_size
    if test_cfg_crop_div_stride != None:
        test_cfg['stride'] = (int(crop_size[0]/test_cfg_crop_div_stride), int(crop_size[1]/test_cfg_crop_div_stride))

    return test_cfg

if __name__ == '__main__':
    # 1. 修改pretrained
    pretrained_pth = './My_Local_Model/open_mmlab/resnet50_v1c.pth'
    pretrained = generate_model_pretrained(pretrained_pth = pretrained_pth)

    # 2. 修改backbone
    depth = 50
    backbone = generate_model_backbone(depth=depth)

    # 3. 修改data_preprocessor
    model_data_preprocessor = 'data_preprocessor'
    model_data_preprocessor = generate_model_data_preprocessor(model_data_preprocessor=model_data_preprocessor)
    
    # 4.5. 修改decode_head、auxiliary_head
    num_classes='36' # 分割数
    decode_head_loss_decode_dict=dict(type='DiceLoss', use_sigmoid=False, loss_weight=1.0) # 直接传入dict
    align_corners=False # 是否需要角对齐
    auxiliary_head_loss_decode_dict=dict(type='DiceLoss', use_sigmoid=False, loss_weight=0.4)

    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)

    # 6. 修改train_cfg
    train_cfg = generate_model_train_cfg()

    # 6. 修改test_cfg
    mode='slide'
    crop_size = (767,767)
    test_cfg_crop_div_stride = 1.5
    test_cfg = generate_model_test_cfg(mode=mode, crop_size=crop_size, test_cfg_crop_div_stride=test_cfg_crop_div_stride)

    # 汇总为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,
    )
    model_ = format_dict(model)
    print(model_)