import torch
# All neural network modules: nn.Linear, nn.Conv2d, BatchNorm, loss functions
import torch.nn as nn
import torch.optim as optim  # For all Optimization algorithms: SGD, Adam, etc
import torch.nn.functional as F  # All functions that don't have any parameters
# Gives easier data management and creates mini batches
from torch.utils.data import DataLoader
# Has standard datasets we can import in a nice way
import torchvision.datasets as datasets
# Transformations we can perform on our datasets
import torchvision.transforms as transforms

VGG_types = {
    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']
}
# Then flatten and 4096 - 4096 - 1000 Linear Layers
# The above dictionary stores all the information about the different VGG architectures


class VGG_net(nn.Module):
    def __init__(self, in_channels=3, num_classes=1000):
        super(VGG_net, self).__init__()
        self.in_channels = in_channels
        self.conv_layers = self.create_conv_layers(VGG_types['VGG16'])

        self.fcs = nn.Sequential(
            nn.Linear(512*7*7, 4096),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(4096, num_classes)
        )

    def forward(self, x):
        x = self.conv_layers(x)
        x = x.reshape(x.shape[0], -1)
        x = self.fcs(x)
        return x

    def create_conv_layers(self, architecture):
        layers = []
        in_channels = self.in_channels

        for x in architecture:
            if type(x) == int:
                out_channels = x

                layers += [nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
                                     kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
                           nn.BatchNorm2d(x),
                           nn.ReLU()]
                in_channels = x
            elif x == 'M':
                layers += [nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))]

        # *layers will unpack all of the content of the list
        return nn.Sequential(*layers)
        # And pytorch will create a sequential model for it

# Testing whether the model gives correct output
model = VGG_net(in_channels=3, num_classes=1000)
x = torch.randn(1, 3, 224, 224)
print(model(x).shape)
