import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision

def EfficientNets(name="efficientnet_b0", num_classes=1860, pretrained=False):

    if name=='efficientnet_b0':
        network = torchvision.models.efficientnet_b0()
    elif name=='efficientnet_b1':
        network = torchvision.models.efficientnet_b1()
    elif name=='efficientnet_b2':
        network = torchvision.models.efficientnet_b2()
    elif name=='efficientnet_b3':
        network = torchvision.models.efficientnet_b3()
    elif name=='efficientnet_b4':
        network = torchvision.models.efficientnet_b4()
    elif name=='efficientnet_b5':
        network = torchvision.models.efficientnet_b5()
    elif name=='efficientnet_b6':
        network = torchvision.models.efficientnet_b6()
    elif name=='efficientnet_b7':
        network = torchvision.models.efficientnet_b7()

    if pretrained:
        network.load_state_dict(torch.load('../models/weights/weights_'+name+'.pth'))

    if network.classifier[1].out_features != num_classes:
        network.classifier[1] = nn.Linear(in_features=network.classifier[1].in_features, out_features=num_classes)

    return network
