import json
MODEL_ZOO = ['resnet50', 'resnet101', 'resnet18', 'resnet50mocov2', 'resnet50byol', 'resnet50simclr',
             'swsl_resnext50_32x4d', 'mobilenetv3_large_100',
             'vit_small_patch16_224mocov3', 'vit_small_patch16_224', 'vit_base_patch16_224mae', 'vit_base_patch16_224',  'vit_tiny_patch16_224', 'vit_large_patch16_224',
             'swin_tiny_patch4_window7_224', 'swin_small_patch4_window7_224', 'swin_base_patch4_window7_224', 'swin_large_patch4_window7_224',
             'regnet_y_16gf', 'regnet_y_3_2gf', 'regnet_y_1_6gf', 'regnet_y_8gf', 'regnet_y_32gf', 'regnet_y_800mf']


MODEL_BLOCKS = {
    'resnet18': ['layer1.0', 'layer1.1', 'layer2.0', 'layer2.1',
                 'layer3.0', 'layer3.1', 'layer4.0', 'layer4.1'],
    'resnet50': ['layer1.0', 'layer1.1', 'layer1.2', 'layer2.0',
                 'layer2.1', 'layer2.2', 'layer2.3', 'layer3.0',
                 'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4',
                 'layer3.5', 'layer4.0', 'layer4.1', 'layer4.2'],
    'resnet50mocov2': ['layer1.0', 'layer1.1', 'layer1.2', 'layer2.0',
                       'layer2.1', 'layer2.2', 'layer2.3', 'layer3.0',
                       'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4',
                       'layer3.5', 'layer4.0', 'layer4.1', 'layer4.2'],
    'resnet50byol': ['layer1.0', 'layer1.1', 'layer1.2', 'layer2.0',
                     'layer2.1', 'layer2.2', 'layer2.3', 'layer3.0',
                     'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4',
                     'layer3.5', 'layer4.0', 'layer4.1', 'layer4.2'],
    'resnet50simclr': ['layer1.0', 'layer1.1', 'layer1.2', 'layer2.0',
                       'layer2.1', 'layer2.2', 'layer2.3', 'layer3.0',
                       'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4',
                       'layer3.5', 'layer4.0', 'layer4.1', 'layer4.2'],
    'resnet101': ['layer1.0', 'layer1.1', 'layer1.2', 'layer2.0',
                  'layer2.1', 'layer2.2', 'layer2.3', 'layer3.0',
                  'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4',
                  'layer3.5', 'layer3.6', 'layer3.7', 'layer3.8',
                  'layer3.9', 'layer3.10', 'layer3.11', 'layer3.12',
                  'layer3.13', 'layer3.14', 'layer3.15', 'layer3.16',
                  'layer3.17', 'layer3.18', 'layer3.19', 'layer3.20',
                  'layer3.21', 'layer3.22', 'layer4.0', 'layer4.1', 'layer4.2'],
    'resnet152': ['layer1.0', 'layer1.1', 'layer1.2', 'layer2.0',
                  'layer2.1', 'layer2.2', 'layer2.3', 'layer2.4',
                  'layer2.5', 'layer2.6', 'layer2.7', 'layer3.0',
                  'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4',
                  'layer3.5', 'layer3.6', 'layer3.7', 'layer3.8',
                  'layer3.9', 'layer3.10', 'layer3.11', 'layer3.12',
                  'layer3.13', 'layer3.14', 'layer3.15', 'layer3.16',
                  'layer3.17', 'layer3.18', 'layer3.19', 'layer3.20',
                  'layer3.21', 'layer3.22', 'layer3.23', 'layer3.24',
                  'layer3.25', 'layer3.26', 'layer3.27', 'layer3.28',
                  'layer3.29', 'layer3.30', 'layer3.31', 'layer3.32',
                  'layer3.33', 'layer3.34', 'layer3.35', 'layer4.0', 'layer4.1', 'layer4.2'],
    'tf_efficientnetv2_l_in21ft1k': [],
    'tf_efficientnetv2_m_in21ft1k': [],
    'tf_efficientnetv2_s_in21ft1k': [],
    'efficientnet_b7': ['features.1.0.block', 'features.1.1.block', 'features.1.2.block',
                        'features.1.3.block', 'features.2.0.block', 'features.2.1.block',
                        'features.2.2.block', 'features.2.3.block', 'features.2.4.block',
                        'features.2.5.block', 'features.2.6.block', 'features.3.0.block',
                        'features.3.1.block', 'features.3.2.block', 'features.3.3.block',
                        'features.3.4.block', 'features.3.5.block', 'features.3.6.block',
                        'features.4.0.block', 'features.4.1.block', 'features.4.2.block',
                        'features.4.3.block', 'features.4.4.block', 'features.4.5.block',
                        'features.4.6.block', 'features.4.7.block', 'features.4.8.block',
                        'features.4.9.block', 'features.5.0.block', 'features.5.1.block',
                        'features.5.2.block', 'features.5.3.block', 'features.5.4.block',
                        'features.5.5.block', 'features.5.6.block', 'features.5.7.block',
                        'features.5.8.block', 'features.5.9.block', 'features.6.0.block',
                        'features.6.1.block', 'features.6.2.block', 'features.6.3.block',
                        'features.6.4.block', 'features.6.5.block', 'features.6.6.block',
                        'features.6.7.block', 'features.6.8.block', 'features.6.9.block',
                        'features.6.10.block', 'features.6.11.block', 'features.6.12.block',
                        'features.7.0.block', 'features.7.1.block', 'features.7.2.block', 'features.7.3.block'],
    'efficientnet_b6': ['features.1.0.block', 'features.1.1.block', 'features.1.2.block',
                        'features.2.0.block', 'features.2.1.block', 'features.2.2.block',
                        'features.2.3.block', 'features.2.4.block', 'features.2.5.block',
                        'features.3.0.block', 'features.3.1.block', 'features.3.2.block',
                        'features.3.3.block', 'features.3.4.block', 'features.3.5.block',
                        'features.4.0.block', 'features.4.1.block', 'features.4.2.block',
                        'features.4.3.block', 'features.4.4.block', 'features.4.5.block',
                        'features.4.6.block', 'features.4.7.block', 'features.5.0.block',
                        'features.5.1.block', 'features.5.2.block', 'features.5.3.block',
                        'features.5.4.block', 'features.5.5.block', 'features.5.6.block',
                        'features.5.7.block', 'features.6.0.block', 'features.6.1.block',
                        'features.6.2.block', 'features.6.3.block', 'features.6.4.block',
                        'features.6.5.block', 'features.6.6.block', 'features.6.7.block',
                        'features.6.8.block', 'features.6.9.block', 'features.6.10.block',
                        'features.7.0.block', 'features.7.1.block', 'features.7.2.block'],
    'efficientnet_b5': ['features.1.0.block', 'features.1.1.block', 'features.1.2.block',
                        'features.2.0.block', 'features.2.1.block', 'features.2.2.block',
                        'features.2.3.block', 'features.2.4.block', 'features.3.0.block',
                        'features.3.1.block', 'features.3.2.block', 'features.3.3.block',
                        'features.3.4.block', 'features.4.0.block', 'features.4.1.block',
                        'features.4.2.block', 'features.4.3.block', 'features.4.4.block',
                        'features.4.5.block', 'features.4.6.block', 'features.5.0.block',
                        'features.5.1.block', 'features.5.2.block', 'features.5.3.block',
                        'features.5.4.block', 'features.5.5.block', 'features.5.6.block',
                        'features.6.0.block', 'features.6.1.block', 'features.6.2.block',
                        'features.6.3.block', 'features.6.4.block', 'features.6.5.block',
                        'features.6.6.block', 'features.6.7.block', 'features.6.8.block',
                        'features.7.0.block', 'features.7.1.block', 'features.7.2.block'],
    'efficientnet_b4': ['features.1.0.block', 'features.1.1.block', 'features.2.0.block',
                        'features.2.1.block', 'features.2.2.block', 'features.2.3.block',
                        'features.3.0.block', 'features.3.1.block', 'features.3.2.block',
                        'features.3.3.block', 'features.4.0.block', 'features.4.1.block',
                        'features.4.2.block', 'features.4.3.block', 'features.4.4.block',
                        'features.4.5.block', 'features.5.0.block', 'features.5.1.block',
                        'features.5.2.block', 'features.5.3.block', 'features.5.4.block',
                        'features.5.5.block', 'features.6.0.block', 'features.6.1.block',
                        'features.6.2.block', 'features.6.3.block', 'features.6.4.block',
                        'features.6.5.block', 'features.6.6.block', 'features.6.7.block',
                        'features.7.0.block', 'features.7.1.block'],
    'efficientnet_b3': ['features.1.0.block', 'features.1.1.block', 'features.2.0.block',
                        'features.2.1.block', 'features.2.2.block', 'features.3.0.block',
                        'features.3.1.block', 'features.3.2.block', 'features.4.0.block',
                        'features.4.1.block', 'features.4.2.block', 'features.4.3.block',
                        'features.4.4.block', 'features.5.0.block', 'features.5.1.block',
                        'features.5.2.block', 'features.5.3.block', 'features.5.4.block',
                        'features.6.0.block', 'features.6.1.block', 'features.6.2.block',
                        'features.6.3.block', 'features.6.4.block', 'features.6.5.block',
                        'features.7.0.block', 'features.7.1.block'],
    'efficientnet_b2': ['features.1.0.block', 'features.1.1.block', 'features.2.0.block',
                        'features.2.1.block', 'features.2.2.block', 'features.3.0.block',
                        'features.3.1.block', 'features.3.2.block', 'features.4.0.block',
                        'features.4.1.block', 'features.4.2.block', 'features.4.3.block',
                        'features.5.0.block', 'features.5.1.block', 'features.5.2.block',
                        'features.5.3.block', 'features.6.0.block', 'features.6.1.block',
                        'features.6.2.block', 'features.6.3.block', 'features.6.4.block',
                        'features.7.0.block', 'features.7.1.block'],
    'efficientnet_b1': ['features.1.0.block', 'features.1.1.block', 'features.2.0.block',
                        'features.2.1.block', 'features.2.2.block', 'features.3.0.block',
                        'features.3.1.block', 'features.3.2.block', 'features.4.0.block',
                        'features.4.1.block', 'features.4.2.block', 'features.4.3.block',
                        'features.5.0.block', 'features.5.1.block', 'features.5.2.block',
                        'features.5.3.block', 'features.6.0.block', 'features.6.1.block',
                        'features.6.2.block', 'features.6.3.block', 'features.6.4.block',
                        'features.7.0.block', 'features.7.1.block'],
    'efficientnet_b0': ['features.1.0.block', 'features.2.0.block', 'features.2.1.block',
                        'features.3.0.block', 'features.3.1.block', 'features.4.0.block',
                        'features.4.1.block', 'features.4.2.block', 'features.5.0.block',
                        'features.5.1.block', 'features.5.2.block', 'features.6.0.block',
                        'features.6.1.block', 'features.6.2.block', 'features.6.3.block', 'features.7.0.block'],
    'mobilenetv3_large_100': ['blocks.0.0', 'blocks.1.0', 'blocks.1.1',
                              'blocks.2.0', 'blocks.2.1', 'blocks.2.2',
                              'blocks.3.0', 'blocks.3.1', 'blocks.3.2',
                              'blocks.3.3', 'blocks.4.0', 'blocks.4.1',
                              'blocks.5.0', 'blocks.5.1', 'blocks.5.2', 'blocks.6.0'],
    'tf_mobilenetv3_large_075': ['blocks.0.0', 'blocks.1.0', 'blocks.1.1',
                                 'blocks.2.0', 'blocks.2.1', 'blocks.2.2',
                                 'blocks.3.0', 'blocks.3.1', 'blocks.3.2',
                                 'blocks.3.3', 'blocks.4.0', 'blocks.4.1',
                                 'blocks.5.0', 'blocks.5.1', 'blocks.5.2', 'blocks.6.0'],
    'tf_mobilenetv3_small_100': ['blocks.0.0', 'blocks.1.0', 'blocks.1.1',
                                 'blocks.2.0', 'blocks.2.1', 'blocks.2.2',
                                 'blocks.3.0', 'blocks.3.1', 'blocks.4.0',
                                 'blocks.4.1', 'blocks.4.2', 'blocks.5.0'],
    'tf_mobilenetv3_small_075': ['blocks.0.0', 'blocks.1.0', 'blocks.1.1',
                                 'blocks.2.0', 'blocks.2.1', 'blocks.2.2',
                                 'blocks.3.0', 'blocks.3.1', 'blocks.4.0',
                                 'blocks.4.1', 'blocks.4.2', 'blocks.5.0'],
    'mobilenetv2_140': ['blocks.0.0', 'blocks.1.0', 'blocks.1.1', 'blocks.2.0',
                        'blocks.2.1', 'blocks.2.2', 'blocks.3.0', 'blocks.3.1',
                        'blocks.3.2', 'blocks.3.3', 'blocks.4.0', 'blocks.4.1',
                        'blocks.4.2', 'blocks.5.0', 'blocks.5.1', 'blocks.5.2', 'blocks.6.0'],
    'mobilenetv2_120d': ['blocks.0.0', 'blocks.1.0', 'blocks.1.1', 'blocks.1.2',
                         'blocks.2.0', 'blocks.2.1', 'blocks.2.2', 'blocks.2.3',
                         'blocks.2.4', 'blocks.3.0', 'blocks.3.1', 'blocks.3.2',
                         'blocks.3.3', 'blocks.3.4', 'blocks.3.5', 'blocks.4.0',
                         'blocks.4.1', 'blocks.4.2', 'blocks.4.3', 'blocks.4.4',
                         'blocks.5.0', 'blocks.5.1', 'blocks.5.2', 'blocks.5.3',
                         'blocks.5.4', 'blocks.6.0'],
    'mobilenetv2_110d': ['blocks.0.0', 'blocks.1.0', 'blocks.1.1', 'blocks.1.2',
                         'blocks.2.0', 'blocks.2.1', 'blocks.2.2', 'blocks.2.3',
                         'blocks.3.0', 'blocks.3.1', 'blocks.3.2', 'blocks.3.3',
                         'blocks.3.4', 'blocks.4.0', 'blocks.4.1', 'blocks.4.2',
                         'blocks.4.3', 'blocks.5.0', 'blocks.5.1', 'blocks.5.2',
                         'blocks.5.3', 'blocks.6.0'],
    'mobilenetv2_100': ['blocks.0.0', 'blocks.1.0', 'blocks.1.1', 'blocks.2.0',
                        'blocks.2.1', 'blocks.2.2', 'blocks.3.0', 'blocks.3.1',
                        'blocks.3.2', 'blocks.3.3', 'blocks.4.0', 'blocks.4.1',
                        'blocks.4.2', 'blocks.5.0', 'blocks.5.1', 'blocks.5.2',
                        'blocks.6.0'],
    'regnet_y_32gf': ['trunk_output.block1.block1-0', 'trunk_output.block1.block1-1', 'trunk_output.block2.block2-0',
                      'trunk_output.block2.block2-1', 'trunk_output.block2.block2-2', 'trunk_output.block2.block2-3',
                      'trunk_output.block2.block2-4', 'trunk_output.block3.block3-0', 'trunk_output.block3.block3-1',
                      'trunk_output.block3.block3-2', 'trunk_output.block3.block3-3',
                      'trunk_output.block3.block3-4', 'trunk_output.block3.block3-5', 'trunk_output.block3.block3-6',
                      'trunk_output.block3.block3-7', 'trunk_output.block3.block3-8', 'trunk_output.block3.block3-9',
                      'trunk_output.block3.block3-10', 'trunk_output.block3.block3-11', 'trunk_output.block4.block4-0'],
    'regnet_y_16gf': ['trunk_output.block1.block1-0', 'trunk_output.block1.block1-1', 'trunk_output.block2.block2-0',
                      'trunk_output.block2.block2-1', 'trunk_output.block2.block2-2', 'trunk_output.block2.block2-3',
                      'trunk_output.block3.block3-0', 'trunk_output.block3.block3-1', 'trunk_output.block3.block3-10',
                      'trunk_output.block3.block3-2', 'trunk_output.block3.block3-3', 'trunk_output.block3.block3-4',
                      'trunk_output.block3.block3-5', 'trunk_output.block3.block3-6', 'trunk_output.block3.block3-7',
                      'trunk_output.block3.block3-8', 'trunk_output.block3.block3-9', 'trunk_output.block4.block4-0'],
    'regnet_y_8gf': ['trunk_output.block1.block1-0', 'trunk_output.block1.block1-1', 'trunk_output.block2.block2-0',
                     'trunk_output.block2.block2-1', 'trunk_output.block2.block2-2', 'trunk_output.block2.block2-3',
                     'trunk_output.block3.block3-0', 'trunk_output.block3.block3-1', 'trunk_output.block3.block3-2',
                     'trunk_output.block3.block3-3', 'trunk_output.block3.block3-4', 'trunk_output.block3.block3-5',
                     'trunk_output.block3.block3-6', 'trunk_output.block3.block3-7', 'trunk_output.block3.block3-8',
                     'trunk_output.block3.block3-9', 'trunk_output.block4.block4-0'],
    'regnet_y_3_2gf': ['trunk_output.block1.block1-0', 'trunk_output.block1.block1-1', 'trunk_output.block2.block2-0',
                       'trunk_output.block2.block2-1', 'trunk_output.block2.block2-2', 'trunk_output.block2.block2-3',
                       'trunk_output.block2.block2-4', 'trunk_output.block3.block3-0', 'trunk_output.block3.block3-1',
                       'trunk_output.block3.block3-2', 'trunk_output.block3.block3-3', 'trunk_output.block3.block3-4',
                       'trunk_output.block3.block3-5', 'trunk_output.block3.block3-6', 'trunk_output.block3.block3-7',
                       'trunk_output.block3.block3-8', 'trunk_output.block3.block3-9', 'trunk_output.block3.block3-10',
                       'trunk_output.block3.block3-11', 'trunk_output.block3.block3-12', 'trunk_output.block4.block4-0'],
    'regnet_y_1_6gf': ['trunk_output.block1.block1-0', 'trunk_output.block1.block1-1', 'trunk_output.block2.block2-0',
                       'trunk_output.block2.block2-1', 'trunk_output.block2.block2-2', 'trunk_output.block2.block2-3',
                       'trunk_output.block2.block2-4', 'trunk_output.block2.block2-5', 'trunk_output.block3.block3-0',
                       'trunk_output.block3.block3-1',  'trunk_output.block3.block3-2', 'trunk_output.block3.block3-3',
                       'trunk_output.block3.block3-4', 'trunk_output.block3.block3-5', 'trunk_output.block3.block3-6',
                       'trunk_output.block3.block3-7', 'trunk_output.block3.block3-8', 'trunk_output.block3.block3-9',
                       'trunk_output.block3.block3-10', 'trunk_output.block3.block3-11', 'trunk_output.block3.block3-12',
                       'trunk_output.block3.block3-13', 'trunk_output.block3.block3-14', 'trunk_output.block3.block3-15',
                       'trunk_output.block3.block3-16', 'trunk_output.block4.block4-0', 'trunk_output.block4.block4-1'],
    'regnet_y_800mf': ['trunk_output.block1.block1-0', 'trunk_output.block2.block2-0', 'trunk_output.block2.block2-1',
                       'trunk_output.block2.block2-2', 'trunk_output.block3.block3-0', 'trunk_output.block3.block3-1',
                       'trunk_output.block3.block3-2', 'trunk_output.block3.block3-3', 'trunk_output.block3.block3-4',
                       'trunk_output.block3.block3-5', 'trunk_output.block3.block3-6', 'trunk_output.block3.block3-7',
                       'trunk_output.block4.block4-0', 'trunk_output.block4.block4-1'],
    # 'swin_base_patch4_window7_224': ['layers.0.blocks.0', 'layers.0.blocks.1', 'layers.1.blocks.0',
    #                                  'layers.1.blocks.1', 'layers.2.blocks.0', 'layers.2.blocks.1',
    #                                  'layers.2.blocks.2', 'layers.2.blocks.3', 'layers.2.blocks.4',
    #                                  'layers.2.blocks.5', 'layers.2.blocks.6', 'layers.2.blocks.7',
    #                                  'layers.2.blocks.8', 'layers.2.blocks.9',  'layers.2.blocks.10',
    #                                  'layers.2.blocks.11', 'layers.2.blocks.12', 'layers.2.blocks.13',
    #                                  'layers.2.blocks.14', 'layers.2.blocks.15', 'layers.2.blocks.16',
    #                                  'layers.2.blocks.17', 'layers.3.blocks.0', 'layers.3.blocks.1'],
    # 'swin_large_patch4_window7_224': ['layers.0.blocks.0', 'layers.0.blocks.1', 'layers.1.blocks.0',
    #                                   'layers.1.blocks.1', 'layers.2.blocks.0', 'layers.2.blocks.1',
    #                                   'layers.2.blocks.2', 'layers.2.blocks.3', 'layers.2.blocks.4',
    #                                   'layers.2.blocks.5', 'layers.2.blocks.6', 'layers.2.blocks.7',
    #                                   'layers.2.blocks.8', 'layers.2.blocks.9', 'layers.2.blocks.10',
    #                                   'layers.2.blocks.11', 'layers.2.blocks.12', 'layers.2.blocks.13',
    #                                   'layers.2.blocks.14', 'layers.2.blocks.15', 'layers.2.blocks.16',
    #                                   'layers.2.blocks.17', 'layers.3.blocks.0', 'layers.3.blocks.1'],
    # 'swin_small_patch4_window7_224': ['layers.0.blocks.0', 'layers.0.blocks.1', 'layers.1.blocks.0',
    #                                   'layers.1.blocks.1', 'layers.2.blocks.0', 'layers.2.blocks.1',
    #                                   'layers.2.blocks.2', 'layers.2.blocks.3', 'layers.2.blocks.4',
    #                                   'layers.2.blocks.5', 'layers.2.blocks.6', 'layers.2.blocks.7',
    #                                   'layers.2.blocks.8', 'layers.2.blocks.9',  'layers.2.blocks.10',
    #                                   'layers.2.blocks.11', 'layers.2.blocks.12', 'layers.2.blocks.13',
    #                                   'layers.2.blocks.14', 'layers.2.blocks.15', 'layers.2.blocks.16',
    #                                   'layers.2.blocks.17', 'layers.3.blocks.0', 'layers.3.blocks.1'],
    # 'swin_tiny_patch4_window7_224': ['layers.0.blocks.0', 'layers.0.blocks.1', 'layers.1.blocks.0', 'layers.1.blocks.1',
    #                                  'layers.2.blocks.0', 'layers.2.blocks.1', 'layers.2.blocks.2', 'layers.2.blocks.3',
    #                                  'layers.2.blocks.4', 'layers.2.blocks.5', 'layers.3.blocks.0', 'layers.3.blocks.1'],
    'swin_base_patch4_window7_224': ['stages.0.blocks.0', 'stages.0.blocks.1', 'stages.0.downsample', 'stages.1.blocks.0',
                                     'stages.1.blocks.1', 'stages.1.downsample', 'stages.2.blocks.0', 'stages.2.blocks.1',
                                     'stages.2.blocks.2', 'stages.2.blocks.3', 'stages.2.blocks.4',
                                     'stages.2.blocks.5', 'stages.2.blocks.6', 'stages.2.blocks.7',
                                     'stages.2.blocks.8', 'stages.2.blocks.9',  'stages.2.blocks.10',
                                     'stages.2.blocks.11', 'stages.2.blocks.12', 'stages.2.blocks.13',
                                     'stages.2.blocks.14', 'stages.2.blocks.15', 'stages.2.blocks.16',
                                     'stages.2.blocks.17', 'stages.2.downsample', 'stages.3.blocks.0', 'stages.3.blocks.1'],
    'swin_large_patch4_window7_224': ['stages.0.blocks.0', 'stages.0.blocks.1', 'stages.0.downsample', 'stages.1.blocks.0',
                                      'stages.1.blocks.1', 'stages.1.downsample', 'stages.2.blocks.0', 'stages.2.blocks.1',
                                      'stages.2.blocks.2', 'stages.2.blocks.3', 'stages.2.blocks.4',
                                      'stages.2.blocks.5', 'stages.2.blocks.6', 'stages.2.blocks.7',
                                      'stages.2.blocks.8', 'stages.2.blocks.9', 'stages.2.blocks.10',
                                      'stages.2.blocks.11', 'stages.2.blocks.12', 'stages.2.blocks.13',
                                      'stages.2.blocks.14', 'stages.2.blocks.15', 'stages.2.blocks.16',
                                      'stages.2.blocks.17', 'stages.2.downsample', 'stages.3.blocks.0', 'stages.3.blocks.1'],
    'swin_small_patch4_window7_224': ['stages.0.blocks.0', 'stages.0.blocks.1', 'stages.0.downsample', 'stages.1.blocks.0',
                                      'stages.1.blocks.1', 'stages.1.downsample', 'stages.2.blocks.0', 'stages.2.blocks.1',
                                      'stages.2.blocks.2', 'stages.2.blocks.3', 'stages.2.blocks.4',
                                      'stages.2.blocks.5', 'stages.2.blocks.6', 'stages.2.blocks.7',
                                      'stages.2.blocks.8', 'stages.2.blocks.9',  'stages.2.blocks.10',
                                      'stages.2.blocks.11', 'stages.2.blocks.12', 'stages.2.blocks.13',
                                      'stages.2.blocks.14', 'stages.2.blocks.15', 'stages.2.blocks.16',
                                      'stages.2.blocks.17', 'stages.2.downsample', 'stages.3.blocks.0', 'stages.3.blocks.1'],
    'swin_tiny_patch4_window7_224': ['stages.0.blocks.0', 'stages.0.blocks.1', 'stages.0.downsample',
                                     'stages.1.blocks.0', 'stages.1.blocks.1', 'stages.1.downsample',
                                     'stages.2.blocks.0', 'stages.2.blocks.1', 'stages.2.blocks.2',
                                     'stages.2.blocks.3', 'stages.2.blocks.4', 'stages.2.blocks.5',
                                     'stages.2.downsample', 'stages.3.blocks.0', 'stages.3.blocks.1'],
    'swsl_resnext101_32x8d': ['layer1.0', 'layer1.1', 'layer1.2', 'layer2.0',
                              'layer2.1', 'layer2.2', 'layer2.3', 'layer3.0',
                              'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4',
                              'layer3.5', 'layer3.6', 'layer3.7', 'layer3.8',
                              'layer3.9', 'layer3.10', 'layer3.11', 'layer3.12',
                              'layer3.13', 'layer3.14', 'layer3.15', 'layer3.16',
                              'layer3.17', 'layer3.18', 'layer3.19', 'layer3.20',
                              'layer3.21', 'layer3.22', 'layer4.0', 'layer4.1', 'layer4.2'],
    'swsl_resnext50_32x4d': ['layer1.0', 'layer1.1', 'layer1.2', 'layer2.0',
                             'layer2.1', 'layer2.2', 'layer2.3', 'layer3.0',
                             'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4',
                             'layer3.5', 'layer4.0', 'layer4.1', 'layer4.2'],
    'resnext50_32x4d': ['layer1.0', 'layer1.1', 'layer1.2', 'layer2.0',
                        'layer2.1', 'layer2.2', 'layer2.3', 'layer3.0',
                        'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4',
                        'layer3.5', 'layer4.0', 'layer4.1', 'layer4.2'],
    'vit_large_patch16_224': [f'blocks.{i}' for i in range(24)],
    'vit_base_patch16_224': [f'blocks.{i}' for i in range(12)],
    'vit_base_patch16_224mae': [f'blocks.{i}' for i in range(12)],
    'vit_base_patch16_224mocov3': [f'blocks.{i}' for i in range(12)],
    'vit_small_patch16_224': [f'blocks.{i}' for i in range(12)],
    'vit_small_patch16_224mocov3': [f'blocks.{i}' for i in range(12)],
    'vit_tiny_patch16_224': [f'blocks.{i}' for i in range(12)],
    'tf_inception_v3': ['Mixed_5b', 'Mixed_5c', 'Mixed_5d', 'Mixed_6a',
                        'Mixed_6b', 'Mixed_6c', 'Mixed_6d', 'Mixed_6e',
                        'Mixed_7a', 'Mixed_7b', 'Mixed_7c'],
    'inception_resnet_v2': [],
    'convnext_tiny': ['stages.0.blocks.0', 'stages.0.blocks.1', 'stages.0.blocks.2',
                      'stages.1.blocks.0', 'stages.1.blocks.1', 'stages.1.blocks.2',
                      'stages.2.blocks.0', 'stages.2.blocks.1', 'stages.2.blocks.2',
                      'stages.2.blocks.3', 'stages.2.blocks.4', 'stages.2.blocks.5',
                      'stages.2.blocks.6', 'stages.2.blocks.7', 'stages.2.blocks.8',
                      'stages.3.blocks.0', 'stages.3.blocks.1', 'stages.3.blocks.2'],
    'convnext_small': ['stages.0.blocks.0', 'stages.0.blocks.1', 'stages.0.blocks.2',
                       'stages.1.blocks.0', 'stages.1.blocks.1', 'stages.1.blocks.2',
                       'stages.2.blocks.0', 'stages.2.blocks.1', 'stages.2.blocks.2', 'stages.2.blocks.3',
                       'stages.2.blocks.4', 'stages.2.blocks.5', 'stages.2.blocks.6', 'stages.2.blocks.7',
                       'stages.2.blocks.8', 'stages.2.blocks.9', 'stages.2.blocks.10', 'stages.3.blocks.0',
                       'stages.2.blocks.11', 'stages.2.blocks.12', 'stages.2.blocks.13', 'stages.2.blocks.14',
                       'stages.2.blocks.15', 'stages.2.blocks.16', 'stages.2.blocks.17', 'stages.2.blocks.18',
                       'stages.2.blocks.19', 'stages.2.blocks.20', 'stages.2.blocks.21', 'stages.2.blocks.22',
                       'stages.2.blocks.23', 'stages.2.blocks.24', 'stages.2.blocks.25', 'stages.2.blocks.26',
                       'stages.3.blocks.1', 'stages.3.blocks.2'],
    'convnext_base': ['stages.0.blocks.0', 'stages.0.blocks.1', 'stages.0.blocks.2', 'stages.1.blocks.0',
                      'stages.1.blocks.1', 'stages.1.blocks.2', 'stages.2.blocks.0',
                      'stages.2.blocks.1',  'stages.2.blocks.2', 'stages.2.blocks.3',
                      'stages.2.blocks.4', 'stages.2.blocks.5', 'stages.2.blocks.6',
                      'stages.2.blocks.7', 'stages.2.blocks.8', 'stages.2.blocks.9',
                      'stages.2.blocks.10', 'stages.2.blocks.11', 'stages.2.blocks.12',
                      'stages.2.blocks.13', 'stages.2.blocks.14', 'stages.2.blocks.15',
                      'stages.2.blocks.16', 'stages.2.blocks.17', 'stages.2.blocks.18',
                      'stages.2.blocks.19', 'stages.2.blocks.20', 'stages.2.blocks.21',
                      'stages.2.blocks.22', 'stages.2.blocks.23', 'stages.2.blocks.24',
                      'stages.2.blocks.25', 'stages.2.blocks.26', 'stages.3.blocks.0',
                      'stages.3.blocks.1', 'stages.3.blocks.2'],
    'convnext_large': ['stages.0.blocks.0', 'stages.0.blocks.1', 'stages.0.blocks.2',
                       'stages.1.blocks.0', 'stages.1.blocks.1', 'stages.1.blocks.2',
                       'stages.2.blocks.0', 'stages.2.blocks.1', 'stages.2.blocks.2',
                       'stages.2.blocks.3', 'stages.2.blocks.4', 'stages.2.blocks.5',
                       'stages.2.blocks.6', 'stages.2.blocks.7', 'stages.2.blocks.8',
                       'stages.2.blocks.9', 'stages.2.blocks.10', 'stages.2.blocks.11',
                       'stages.2.blocks.12', 'stages.2.blocks.13', 'stages.2.blocks.14',
                       'stages.2.blocks.15', 'stages.2.blocks.16', 'stages.2.blocks.17',
                       'stages.2.blocks.18', 'stages.2.blocks.19', 'stages.2.blocks.20',
                       'stages.2.blocks.21', 'stages.2.blocks.22', 'stages.2.blocks.23',
                       'stages.2.blocks.24', 'stages.2.blocks.25', 'stages.2.blocks.26',
                       'stages.3.blocks.0', 'stages.3.blocks.1', 'stages.3.blocks.2'],
    'repvgg_b3': ['stages.0.0', 'stages.0.1', 'stages.0.2', 'stages.0.3',
                  'stages.1.0', 'stages.1.1', 'stages.1.2', 'stages.1.3',
                  'stages.1.4', 'stages.1.5', 'stages.2.0', 'stages.2.1',
                  'stages.2.2', 'stages.2.3', 'stages.2.4', 'stages.2.5',
                  'stages.2.6', 'stages.2.7', 'stages.2.8', 'stages.2.9',
                  'stages.2.10', 'stages.2.11', 'stages.2.12', 'stages.2.13',
                  'stages.2.14', 'stages.2.15', 'stages.3.0'],
    'repvgg_b2': ['stages.0.0', 'stages.0.1', 'stages.0.2', 'stages.0.3',
                  'stages.1.0', 'stages.1.1', 'stages.1.2', 'stages.1.3',
                  'stages.1.4', 'stages.1.5', 'stages.2.0', 'stages.2.1',
                  'stages.2.2', 'stages.2.3', 'stages.2.4', 'stages.2.5',
                  'stages.2.6', 'stages.2.7', 'stages.2.8', 'stages.2.9',
                  'stages.2.10', 'stages.2.11', 'stages.2.12', 'stages.2.13',
                  'stages.2.14', 'stages.2.15', 'stages.3.0'],
    'repvgg_b1': ['stages.0.0', 'stages.0.1', 'stages.0.2', 'stages.0.3',
                  'stages.1.0', 'stages.1.1', 'stages.1.2', 'stages.1.3',
                  'stages.1.4', 'stages.1.5', 'stages.2.0', 'stages.2.1',
                  'stages.2.2', 'stages.2.3', 'stages.2.4', 'stages.2.5',
                  'stages.2.6', 'stages.2.7', 'stages.2.8', 'stages.2.9',
                  'stages.2.10', 'stages.2.11', 'stages.2.12', 'stages.2.13',
                  'stages.2.14', 'stages.2.15', 'stages.3.0'],
    'repvgg_b0': ['stages.0.0', 'stages.0.1', 'stages.0.2', 'stages.0.3',
                  'stages.1.0', 'stages.1.1', 'stages.1.2', 'stages.1.3',
                  'stages.1.4', 'stages.1.5', 'stages.2.0', 'stages.2.1',
                  'stages.2.2', 'stages.2.3', 'stages.2.4', 'stages.2.5',
                  'stages.2.6', 'stages.2.7', 'stages.2.8', 'stages.2.9',
                  'stages.2.10', 'stages.2.11', 'stages.2.12', 'stages.2.13',
                  'stages.2.14', 'stages.2.15', 'stages.3.0'],
}

# ResNet 3 + efficientnet 8 + mobilenetv3 2 + regnet_y 4 + swin 4 + VIT 4 + resnext 2

MODEL_PRINT = {
    'resnet18': 'ResNet18',
    'resnet50': 'ResNet50',
    'resnet50mocov2': 'ResNet50 MoCov2',
    'resnet50byol': 'ResNet50 BYOL',
    'resnet50simclr': 'ResNet50 SimCLR',
    'resnet101': 'ResNet101',
    'resnet152': 'ResNet152',
    'tf_efficientnetv2_l_in21ft1k': [],
    'tf_efficientnetv2_m_in21ft1k': [],
    'tf_efficientnetv2_s_in21ft1k': [],
    'efficientnet_b7': 'EffiNet b7',
    'efficientnet_b6': 'EffiNet b6',
    'efficientnet_b5': 'EffiNet b5',
    'efficientnet_b4': 'EffiNet b4',
    'efficientnet_b3': 'EffiNet b3',
    'efficientnet_b2': 'EffiNet b2',
    'efficientnet_b1': 'EffiNet b1',
    'efficientnet_b0': 'EffiNet b0',
    'mobilenetv3_large_100': 'MobileNetv3 Large 1.0',
    'mobilenetv3_large_075': 'MobileNetv3 Large 0.75',
    'tf_mobilenetv3_large_075': 'MobileNetv3 Large 0.75',
    'tf_mobilenetv3_small_100': 'MobileNetv3 Small 1.0',
    'tf_mobilenetv3_small_075': 'MobileNetv3 Small 0.75',
    'regnet_y_16gf': 'RegNetY-16GF',
    'regnet_y_8gf': 'RegNetY-8GF',
    'regnet_y_32gf': 'RegNetY-32GF',
    'regnet_y_3_2gf': 'RegNetY-3.2GF',
    'regnet_y_1_6gf': 'RegNetY-1.6GF',
    'regnet_y_800mf': 'RegNetY-800MF',
    'swin_base_patch4_window7_224': 'Swin-B',
    'swin_large_patch4_window7_224': 'Swin-L',
    'swin_small_patch4_window7_224': 'Swin-S',
    'swin_tiny_patch4_window7_224': 'Swin-T',
    'swsl_resnext101_32x8d': ['layer1.0', 'layer1.1', 'layer1.2', 'layer2.0',
                              'layer2.1', 'layer2.2', 'layer2.3', 'layer3.0',
                              'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4',
                              'layer3.5', 'layer3.6', 'layer3.7', 'layer3.8',
                              'layer3.9', 'layer3.10', 'layer3.11', 'layer3.12',
                              'layer3.13', 'layer3.14', 'layer3.15', 'layer3.16',
                              'layer3.17', 'layer3.18', 'layer3.19', 'layer3.20',
                              'layer3.21', 'layer3.22', 'layer4.0', 'layer4.1', 'layer4.2'],
    'swsl_resnext50_32x4d': 'ResNeXt50 32x4d',
    'resnext50_32x4d': ['layer1.0', 'layer1.1', 'layer1.2', 'layer2.0',
                        'layer2.1', 'layer2.2', 'layer2.3', 'layer3.0',
                        'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4',
                        'layer3.5', 'layer4.0', 'layer4.1', 'layer4.2'],
    'vit_small_patch16_224mocov3': 'ViT-S MoCov3',
    'vit_base_patch16_224mocov3': 'ViT-B MoCov3',
    'vit_base_patch16_224mae': 'ViT-B MAE',
    'vit_large_patch16_224': 'ViT-L',
    'vit_base_patch16_224': 'ViT-B',
    'vit_small_patch16_224': 'ViT-S',
    'vit_tiny_patch16_224': 'ViT-T',
}


MODEL_STATS = {
    'resnet18': dict(arch='resnet18', top1=69.758, param=11.69, backend='pytorch', type='cnn'),
    'resnet50': dict(arch='resnet50', top1=76.130, param=25.56, backend='pytorch', type='cnn'),
    'resnet50mocov2': dict(arch='resnet50', top1=71.1, param=25.56, backend='pytorch', type='cnn'),
    'resnet50byol': dict(arch='resnet50', top1=67.68, param=25.56, backend='pytorch', type='cnn'),
    'resnet50simclr': dict(arch='resnet50', top1=58.92, param=25.56, backend='pytorch', type='cnn'),
    'resnet101': dict(arch='resnet101', top1=77.374, param=44.55, backend='pytorch', type='cnn'),
    'resnet152': dict(arch='resnet152', top1=78.312, param=60.19, backend='pytorch', type='cnn'),
    'efficientnet_b7': dict(arch='efficientnet_b7', top1=84.122, param=66.35, backend='pytorch', type='cnn'),
    'efficientnet_b6': dict(arch='efficientnet_b6', top1=84.008, param=43.04, backend='pytorch', type='cnn'),
    'efficientnet_b5': dict(arch='efficientnet_b5', top1=83.444, param=30.39, backend='pytorch', type='cnn'),
    'efficientnet_b4': dict(arch='efficientnet_b4', top1=83.384, param=19.34, backend='pytorch', type='cnn'),
    'efficientnet_b3': dict(arch='efficientnet_b3', top1=82.008, param=12.23, backend='pytorch', type='cnn'),
    'efficientnet_b2': dict(arch='efficientnet_b2', top1=80.608, param=9.11, backend='pytorch', type='cnn'),
    'efficientnet_b1': dict(arch='efficientnet_b1', top1=78.642, param=7.79, backend='pytorch', type='cnn'),
    'efficientnet_b0': dict(arch='efficientnet_b0', top1=77.692, param=5.29, backend='pytorch', type='cnn'),
    'mobilenetv3_large_100': dict(arch='mobilenetv3_large_100', top1=75.762, param=5.48, backend='mytimm', type='cnn'),
    'mobilenetv3_large_075': dict(arch='mobilenetv3_large_075', top1=73.440, param=3.99, backend='mytimm', type='cnn'),
    'tf_mobilenetv3_large_075': dict(arch='mobilenetv3_large_075', top1=73.440, param=3.99, backend='mytimm', type='cnn'),
    'tf_mobilenetv3_small_100': dict(arch='mobilenetv3_small_100', top1=67.924, param=2.54, backend='mytimm', type='cnn'),
    'tf_mobilenetv3_small_075': dict(arch='mobilenetv3_small_075', top1=65.714, param=2.04, backend='mytimm', type='cnn'),
    'regnet_y_32gf': dict(arch='regnet_y_32gf', top1=80.808, param=145.05, backend='pytorch', type='cnn'),
    'regnet_y_16gf': dict(arch='regnet_y_16gf', top1=80.424, param=83.59, backend='pytorch', type='cnn'),
    'regnet_y_8gf': dict(arch='regnet_y_8gf', top1=80.032, param=39.18, backend='pytorch', type='cnn'),
    'regnet_y_3_2gf': dict(arch='regnet_y_3_2gf', top1=78.948, param=19.44, backend='pytorch', type='cnn'),
    'regnet_y_1_6gf': dict(arch='regnet_y_1_6gf', top1=77.950, param=11.20, backend='pytorch', type='cnn'),
    'regnet_y_800mf': dict(arch='regnet_y_800mf', top1=76.312, param=6.26, backend='pytorch', type='cnn'),
    'swin_base_patch4_window7_224': dict(arch='swin_base_patch4_window7_224', top1=85.250, param=87.77, backend='mmcv', 
                                        cfg='configs/_base_/models/swin_transformer/base_224.py', 
                                        type='swin',
                                        load_from = 'https://download.openmmlab.com/mmclassification/v0/swin-transformer/convert/swin_base_patch4_window7_224_22kto1k-f967f799.pth'),
    'swin_large_patch4_window7_224': dict(arch='swin_large_patch4_window7_224', top1=86.318, param=196.53, backend='mmcv', 
                                        cfg='configs/_base_/models/swin_transformer/large_224.py', 
                                        type='swin',
                                        load_from = 'https://download.openmmlab.com/mmclassification/v0/swin-transformer/convert/swin_large_patch4_window7_224_22kto1k-5f0996db.pth'),
    'swin_small_patch4_window7_224': dict(arch='swin_small_patch4_window7_224', top1=83.214, param=49.61, backend='mmcv', 
                                        cfg='configs/_base_/models/swin_transformer/small_224.py', 
                                        type='swin',
                                        load_from='https://download.openmmlab.com/mmclassification/v0/swin-transformer/convert/swin_small_patch4_window7_224-cc7a01c9.pth'),
    'swin_tiny_patch4_window7_224': dict(arch='swin_tiny_patch4_window7_224', top1=81.374, param=28.29, backend='mmcv', 
                                        cfg='configs/_base_/models/swin_transformer/tiny_224.py', 
                                        type='swin',
                                        load_from = 'https://download.openmmlab.com/mmclassification/v0/swin-transformer/swin_tiny_224_b16x64_300e_imagenet_20210616_090925-66df6be6.pth'),
    'swsl_resnext101_32x8d': dict(arch='swsl_resnext101_32x8d', top1=84.290, param=88.79, backend='mytimm', type='cnn'),
    'swsl_resnext50_32x4d': dict(arch='swsl_resnext50_32x4d', top1=82.174, param=25.03, backend='mytimm', type='cnn'),
    'vit_small_patch16_224mocov3': dict(arch='vit_small_patch16_224', top1=73.2, param=22.05, backend='mytimm', type='vit'),
    'vit_base_patch16_224mocov3': dict(arch='vit_base_patch16_224', top1=68.0, param=86.57, backend='mytimm', type='vit'),
    'vit_base_patch16_224mae': dict(arch='vit_base_patch16_224', top1=68.0, param=86.57, backend='mytimm', type='vit'),
    'vit_large_patch16_224': dict(arch='vit_large_patch16_224', top1=85.844, param=304.33, backend='mytimm', type='vit'),
    'vit_base_patch16_224': dict(arch='vit_base_patch16_224', top1=84.526, param=86.57, backend='mytimm', type='vit'),
    'vit_small_patch16_224': dict(arch='vit_small_patch16_224', top1=81.396, param=22.05, backend='mytimm', type='vit'),
    'vit_tiny_patch16_224': dict(arch='vit_tiny_patch16_224', top1=75.456, param=5.72, backend='mytimm', type='vit'),
}


try:
    with open('tools/MODEL_INOUT_SHAPE.json') as json_file:
        MODEL_INOUT_SHAPE = json.load(json_file)
except:
    MODEL_INOUT_SHAPE = None
    print('tools/MODEL_INOUT_SHAPE.json not loaded')
