import torch.nn as nn

from efficientnet_pytorch import EfficientNet as efficientnet

from torchvision import models

class HeartNet(nn.Module):

    def __init__(self, num_classes=7):

        super(HeartNet, self).__init__()

        self.features = nn.Sequential(

            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),

            nn.ELU(inplace=True),

            nn.BatchNorm2d(64, eps=0.001),

            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),

            nn.ELU(inplace=True),

            nn.BatchNorm2d(64, eps=0.001),

            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),

            nn.ELU(inplace=True),

            nn.BatchNorm2d(128, eps=0.001),

            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),

            nn.ELU(inplace=True),

            nn.BatchNorm2d(128, eps=0.001),

            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),

            nn.ELU(inplace=True),

            nn.BatchNorm2d(256, eps=0.001),

            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),

            nn.ELU(inplace=True),

            nn.BatchNorm2d(256, eps=0.001),

            nn.MaxPool2d(kernel_size=2, stride=2),

        )

        self.classifier = nn.Sequential(

            nn.Linear(16 * 16 * 256, 2048),

            nn.ELU(inplace=True),

            nn.BatchNorm1d(2048, eps=0.001),

            nn.Dropout(0.5),

            nn.Linear(2048, num_classes),

        )

    def forward(self, x):

        x = self.features(x)

        x = x.view(x.size(0), 16 * 16 * 256)

        x = self.classifier(x)

        return x

class MobileNetV2(models.MobileNetV2):

    def __init__(self, num_classes=8):

        super().__init__(num_classes=num_classes)

class AlexNet(models.AlexNet):

    def __init__(self, num_classes=8):

        super().__init__(num_classes=num_classes)

def VGG16(num_classes=8):

    model = models.vgg16_bn()

    model.classifier[-1] = nn.Linear(model.classifier[-1].in_features, num_classes)

    return model

def ResNet18(num_classes=8):

    model = models.resnet18()

    model.fc = nn.Linear(model.fc.in_features, num_classes)

    return model

def ResNet34(num_classes=8):

    model = models.resnet34()

    model.fc = nn.Linear(model.fc.in_features, num_classes)

    return model

def ShuffleNet(num_classes=8):

    model = models.shufflenet_v2_x1_0()

    model.fc = nn.Linear(model.fc.in_features, num_classes)

    return model

def EfficientNetB4(num_classes=8):

    model = efficientnet.from_name("efficientnet-b4")

    model._fc = nn.Linear(model._fc.in_features, num_classes)

    return model