import os, sys

import torch

import torchvision

from torch import nn

def loadResnetBackbone(arch, torch_pretrain= None, freeze= False):

    ## pretrain setting

    if torch_pretrain in ["DEFAULT", "IMAGENET-1K"]:

        torch_pretrain = "DEFAULT" #"IMAGENET1K_V2"

    elif torch_pretrain in [None, "NONE", "none"]:

        torch_pretrain = None

    else:

        raise ValueError("Unknown pretrain weight type requested ", torch_pretrain )

    print("Torch Pretrain Set to ...", torch_pretrain)

    ## Model loading

    if arch == 'resnet18':

        backbone = torchvision.models.resnet18(zero_init_residual=True,

                                weights=torch_pretrain)

        outfeat_size = 512

    elif arch == 'resnet34':

        backbone = torchvision.models.resnet34(zero_init_residual=True,

                            weights=torch_pretrain)

        outfeat_size = 512

    elif arch == 'resnet50':

        backbone = torchvision.models.resnet50(zero_init_residual=True,

                            weights=torch_pretrain)

        outfeat_size = 2048

    elif arch == 'resnet101':

        backbone = torchvision.models.resnet101(zero_init_residual=True,

                            weights=torch_pretrain)

        outfeat_size = 2048

    elif arch == 'resnet152':

        backbone = torchvision.models.resnet152(zero_init_residual=True,

                            weights=torch_pretrain)

        outfeat_size = 2048

    else:

        raise ValueError(f"Unknown Model Implementation called in {os.path.basename(__file__)}")

    backbone.fc = nn.Identity() #remove fc of default arch

    # freeze model

    if freeze:

        print("Freezing Resnet weights ...")

        for param in backbone.parameters():

            param.requires_grad = False

    return backbone, outfeat_size