import torch
import torch.nn as nn


VGG_CONFIG = {
    "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"]
}


class VGGCustom(nn.Module):
    def __init__(self,
                 architecture:str="VGG16",
                 num_channels:int = 3,
                 num_classes:int=10,
                 batch_normalization=False,
                 dropout:float=0.5) -> None:

        super().__init__()

        if architecture not in VGG_CONFIG:
            raise ValueError(f"Architecture {architecture} is invalid. ",
                             f"Available architectures are :{list(VGG_CONFIG.keys())}")


        self.batch_normalization = batch_normalization

        self.features = self._make_layers(
            VGG_CONFIG[architecture],
            in_channels=num_channels
        )

        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))

        self.classifier = nn.Sequential(
            nn.Linear(512*7*7, 4096),
            nn.ReLU(True),
            nn.Dropout(p=dropout),

            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(p=dropout),

            nn.Linear(4096, num_classes)
        )

    def _make_layers(self, cfg, in_channels):
        layers = []

        for layer in cfg:
            if layer == 'M':
                layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
            else:

                layers.append(nn.Conv2d(in_channels=in_channels,
                                        out_channels=layer,
                                        kernel_size=3,
                                        stride=1,
                                        padding=1))

                if self.batch_normalization:
                    layers.append(nn.BatchNorm2d(layer))

                layers.append(nn.ReLU(True))

                in_channels = layer

        return nn.Sequential(*layers)

    def forward(self, image):
        x = self.features(image)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)

        return x