# Code from: https://github.com/thu-ml/adversarial_training_imagenet
# @article{liu2023comprehensive,
#   title={A Comprehensive Study on Robustness of Image Classification Models: Benchmarking and Rethinking},
#   author={Liu, Chang and Dong, Yinpeng and Xiang, Wenzhao and Yang, Xiao and Su, Hang and Zhu, Jun and Chen, Yuefeng and He, Yuan and Xue, Hui and Zheng, Shibao},
#   journal={arXiv preprint arXiv:2302.14301},
#   year={2023}
# }

from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD

MEAN_ZERO = (0.0, 0.0, 0.0)
STA_ONE = (1.0, 1.0, 1.0)

url_googleid = {
    "ARES_Resnet50_MOCO.pth": "1dWVD-LvFpd8XJ39psJz4qymv2DBVE4CG",
    "ARES_Resnet50_AT.pth": "1lNjFX_YT0BbIu1OwxHfFDrB_ZUIIqzQ0",
    "ARES_ResNet50_RB1.pth": "1PrdP2MqXbfea__F0k9Tnv2ZQOxACPrMq",
    "ARES_ResNet50_RB2.pth": "1_xfv4jpDdavwUg92xch-G7kAhvbnFZdY",
    "ARES_Resnet101_AT.pth": "1hgqZ3egXsXvo9AIfKGVKqQSmlIKOhMaj",
    "ARES_Resnet152_AT.pth": "1Ti3IeikzCruy8KFSaRBI9b0pQO5I01yn",
    "ARES_Resnet152_FD.pth": "1gDSkuO-z5cqtqd_K--AiaFDByyuGOzs7",
    "ARES_Wide_Resnet50_2_AT.pth": "1Sd-8qdsqJPnGHiy1Q6nluQQmgdXFzqCl",
    "ARES_Wide_Resnet50_2_RB.pth": "1_Di7BEKzPflW9OL2ZhGE4ud55N-V9V5H",
    "ARES_Swin_small_patch4_window7_224_21K.pth": "1KrRAP253r7fXRc99sJYeRPWEtLxRzw9C",
    "ARES_Swin_small_patch4_window7_224_AT.pth": "1YzimzhbUt5X3YsgRIXKPHG8jrIVAhcSD",
    "ARES_Swin_base_patch4_window7_224_Normal.pth": "1DC98d0sMfxqmvrAFOWYEXdiPRYhRpqoz",
    "ARES_Swin_base_patch4_window7_224_AT.pth": "1e_fHf4YuWIxWVmTJNgwELZk7_taRk5bq",
    "ARES_Swin_large_patch4_window7_224_AT.pth": "1Ebag3bFZT86kt5qXJk39m06YfBBg_uPR",
    "ARES_ViT_small_patch16_224_Normal.pth": "1MhabzOSpnXRrcAWKD5n6KkhwZAEEc06T",
    "ARES_ViT_small_patch16_224_AT.pth": "1_4dhvkGmFT8OJLk1pijzHs2LL89HG9oR",
    "ARES_ViT_base_patch16_224_Normal.pth": "18Muu7d0j7z6PeqGWzAp69IgsoODIk4fF",
    "ARES_ViT_base_patch16_224_AT.pth": "1tTMfitpnrP6SyZ6zyxXU3X5IVpllBHCG",
    "ARES_ViT_base_patch16_224_MAE.pth": "1wUXTn2WmhFPjGjo-FxZnkOKJwEbQOy6i",
    "ARES_ViT_large_patch16_224_Normal.pth": "1Avm_-yHZdW0LN2S2Byc2flqi3-NrAdES",
    "ARES_ViT_large_patch16_224_MAE.pth": "18V7KH4fJOwekxFzJ26Z7dZV2D8HuVkWJ",
    "ARES_T2T_ViT_14_Normal.pth": "1a6bg-91nESF12pJEMhVIiJcs8vNO9yWi",
    "ARES_T2T_ViT_19_Normal.pth": "17XXwWt9M0d8c5yGf32aGMXDYY-94bK-9",
    "ARES_T2T_ViT_24_Normal.pth": "1VMIFARp5zw_E2XkUH-rIaNUcKXYqjyjS",
    "ARES_XciT_small_12_p16_224_Normal.pth": "1gE-NV4q8fAAJMCnd8oXAC9m1TnT8k6iw",
    "ARES_XciT_medium_12_p16_224_Normal.pth": "1GlQu2SKgP5sHLqR9CEaObTEjgRSv3itV",
    "ARES_XciT_large_12_p16_224_Normal.pth": "1x_Gmhq4-1G8Sdm9F1gEoOefxbNn-Aq3X",
    "ARES_ConvNext_Small_AT.pth": "1i7wcEt0pewMtww9ZoAn_z0tUKP5ujKdp",
    "ARES_ConvNext_Base_AT.pth": "1NnlLG1uGnuKPBK5DdlRKoyVHGRM7pwm-",
    "ARES_ConvNext_Large_AT.pth": "1Lk8_l265xiInWTIxgiJqKI-knqHEVF6i",
}


def get_url(ckpt_name):
    url = "https://drive.google.com/uc?export=download&id=" + url_googleid[ckpt_name]

    return url


model_zoo = {
    "resnet50_normal": {
        "model": "tv_resnet50",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "resnet50_moco": {
        "model": "tv_resnet50",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Resnet50_MOCO.pth"),
    },
    "resnet50_at": {
        "model": "resnet50",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Resnet50_AT.pth"),
    },
    "resnet50_sota1": {
        "model": "resnet50_rl",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ResNet50_RB1.pth"),
    },
    "resnet50_sota2": {
        "model": "resnet50_rl",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ResNet50_RB2.pth"),
    },
    "resnet101_normal": {
        "model": "tv_resnet101",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "resnet101_at": {
        "model": "resnet101",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Resnet101_AT.pth"),
    },
    "resnet152_normal": {
        "model": "tv_resnet152",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 5,
        "url": "",
    },
    "resnet152_at": {
        "model": "resnet152",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": True,
        "batch_size": 5,
        "url": get_url("ARES_Resnet152_AT.pth"),
    },
    "resnet152_fd": {
        "model": "resnet152_fd",
        "mean": MEAN_ZERO,
        "std": STA_ONE,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 5,
        "url": get_url("ARES_Resnet152_FD.pth"),
    },
    "wresnet50_normal": {
        "model": "wide_resnet50_2",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "wresnet50_at": {
        "model": "wide_resnet50_2",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": True,
        "batch_size": 25,
        "url": get_url("ARES_Wide_Resnet50_2_AT.pth"),
    },
    "wresnet50_sota": {
        "model": "wide_resnet50_2_rl",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Wide_Resnet50_2_RB.pth"),
    },
    "vgg13_normal": {
        "model": "vgg13",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "vgg16_normal": {
        "model": "vgg16",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "vgg19_normal": {
        "model": "vgg19",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "densenet121_normal": {
        "model": "tv_densenet121",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "densenet161_normal": {
        "model": "densenet161",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "densenet201_normal": {
        "model": "densenet201",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "convs_normal": {
        "model": "convnext_small",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "convs_21k": {
        "model": "convnext_small_in22ft1k",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "convb_normal": {
        "model": "convnext_base",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "convb_21k": {
        "model": "convnext_base_in22ft1k",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "convl_normal": {
        "model": "convnext_large",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 5,
        "url": "",
    },
    "convl_21k": {
        "model": "convnext_large_in22ft1k",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 5,
        "url": "",
    },
    "swins_normal": {
        "model": "swin_small_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "swins_21k": {
        "model": "swin_small_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Swin_small_patch4_window7_224_21K.pth"),
    },
    "swins_at": {
        "model": "swin_small_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Swin_small_patch4_window7_224_AT.pth"),
    },
    "swinb_normal": {
        "model": "swin_base_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Swin_base_patch4_window7_224_Normal.pth"),
    },
    "swinb_21k": {
        "model": "swin_base_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth",
    },
    "swinb_at": {
        "model": "swin_base_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_Swin_base_patch4_window7_224_AT.pth"),
    },
    "swinl_21k": {
        "model": "swin_large_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 5,
        "url": "",
    },
    "swinl_at": {
        "model": "swin_large_patch4_window7_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 5,
        "url": get_url("ARES_Swin_large_patch4_window7_224_AT.pth"),
    },
    "vits_normal": {
        "model": "vit_small_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ViT_small_patch16_224_Normal.pth"),
    },
    "vits_21k": {
        "model": "vit_small_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "vits_at": {
        "model": "vit_small_patch16_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ViT_small_patch16_224_AT.pth"),
    },
    "vitb_normal": {
        "model": "vit_base_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ViT_base_patch16_224_Normal.pth"),
    },
    "vitb_21k": {
        "model": "vit_base_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "vitb_at": {
        "model": "vit_base_patch16_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ViT_base_patch16_224_AT.pth"),
    },
    "vitb_mae": {
        "model": "vit_base_patch16",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ViT_base_patch16_224_MAE.pth"),
    },
    "vitl_normal": {
        "model": "vit_large_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 5,
        "url": get_url("ARES_ViT_large_patch16_224_Normal.pth"),
    },
    "vitl_21k": {
        "model": "vit_large_patch16_224",
        "mean": IMAGENET_INCEPTION_MEAN,
        "std": IMAGENET_INCEPTION_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 5,
        "url": "",
    },
    "vitl_mae": {
        "model": "vit_large_patch16",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 5,
        "url": get_url("ARES_ViT_large_patch16_224_MAE.pth"),
    },
    "t2t14_normal": {
        "model": "t2t_vit_t_14",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_T2T_ViT_14_Normal.pth"),
    },
    "t2t19_normal": {
        "model": "t2t_vit_t_19",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_T2T_ViT_19_Normal.pth"),
    },
    "t2t24_normal": {
        "model": "t2t_vit_t_24",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_T2T_ViT_24_Normal.pth"),
    },
    "xcits_normal": {
        "model": "xcit_small_24_p16_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "xcits_sota": {
        "model": "xcit_small_12_p16_224",
        "mean": MEAN_ZERO,
        "std": STA_ONE,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_XciT_small_12_p16_224_Normal.pth"),
    },
    "xcitm_normal": {
        "model": "xcit_medium_24_p16_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "xcitm_sota": {
        "model": "xcit_medium_12_p16_224",
        "mean": MEAN_ZERO,
        "std": STA_ONE,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_XciT_medium_12_p16_224_Normal.pth"),
    },
    "xcitl_normal": {
        "model": "xcit_large_24_p16_224",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": True,
        "act_gelu": False,
        "batch_size": 25,
        "url": "",
    },
    "xcitl_sota": {
        "model": "xcit_large_12_p16_224",
        "mean": MEAN_ZERO,
        "std": STA_ONE,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_XciT_large_12_p16_224_Normal.pth"),
    },
    "convnexts_at": {
        "model": "convnext_small",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ConvNext_Small_AT.pth"),
    },
    "convnextb_at": {
        "model": "convnext_base",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ConvNext_Base_AT.pth"),
    },
    "convnextl_at": {
        "model": "convnext_large",
        "mean": IMAGENET_DEFAULT_MEAN,
        "std": IMAGENET_DEFAULT_STD,
        "pretrained": False,
        "act_gelu": False,
        "batch_size": 25,
        "url": get_url("ARES_ConvNext_Large_AT.pth"),
    },
}
