#Filename:	vgg.py
#Author:	Wang Yongjie
#Email:		yongjie.wang@ntu.edu.sg
#Date:		Min 31 Mei 2020 11:15:47  WIB

"""
Define VGG models with getBiases() and getFeatures() method
For correct computation of full gradient, do not use inplace operation inside the model

"""

import torch
import torch.nn as nn
import torch.utils.model_zoo as model_zoo

cfg = {
    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}

model_urls = {
    'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
    'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth',
    'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
    'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth',
    'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth',
    'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth',
    'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
    'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
}

class VGG(nn.Module):

    def __init__(self, vgg_name, batch_norm = False, num_classes = 1000, init_weights = True):
        super(VGG, self).__init__()
        self.features = self.make_layers(cfg[vgg_name], batch_norm = batch_norm)
        self.name = vgg_name
        self.bn = batch_norm

        self.classifier = nn.Sequential(
                nn.Linear(512 * 7 * 7, 4096),
                nn.ReLU(False),
                nn.Dropout(),
                nn.Linear(4096, 4096),
                nn.ReLU(False),
                nn.Dropout(),
                nn.Linear(4096, num_classes),
                )

        if init_weights == True:
            self._init_weights()

        self.get_biases = False
        self.get_features = False

        self.biases = []
        self.feature_list = []

    def getBiases(self):
        """
        Returns the explicit biases arising from BatchNorm of Convolutional layers
        """
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.get_biases = True 
        self.biases = [0]

        x = torch.zeros(1, 3, 224, 224).to(device)
        _ = self.forward(x)
        self.get_biases = False

        return self.biases

    def getFeatures(self, x):
        """
        Return features at every layer before the ReLU
        """
        self.get_features = True
        self.feature_list = [x]

        x = self.forward(x)
        self.get_features = False
        return x, self.feature_list

    def _classify(self, x):

        for m in self.classifier:
            x = m(x)
            if isinstance(m, nn.Linear):
                if self.get_biases:
                    self.biases.append(m.bias)
                if self.get_features:
                    self.feature_list.append(x)
        
        return x

    def forward(self, x):
        x = self.organize_features(x)
        x = x.view(x.size(0), -1)
        x = self._classify(x)
        return x

    def organize_features(self, x):
        in_channels = 3
        count = 0

        for i in cfg[self.name]:
            if (i == "M"):
                x = self.features[count](x)
            else:
                if self.get_biases:
                    input_bias = torch.zeros(x.size()).to(x.device) # Zero tensor
                    input_bias, _ = self._linear_block(input_bias, count)
                    self.biases.append(input_bias.detach())

                x, count = self._linear_block(x, count)

                if self.get_features:
                    self.feature_list.append(x)

                x = self.features[count](x) # ReLU

            count = count + 1

        return x

    def _linear_block(self, x, count):

        if self.bn: # Conv + BN
            x = self.features[count](x)
            count = count + 1
        x = self.features[count](x) # Conv
        count = count + 1
        return x, count

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode = 'fan_out', nonlinearity = "relu")
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0.01)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1.)
                nn.init.constant_(m.bias, 0.01)
            elif isinstance(m. nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

    def make_layers(self, cfg, batch_norm = True):

        layers = []
        in_channels = 3
        index = 0

        for v in cfg:
            if v == 'M':
                layers += [nn.MaxPool2d(kernel_size = 2, stride = 2)]
            else:
                conv2d = nn.Conv2d(in_channels, v, kernel_size = 3, padding = 1)
                if batch_norm:
                    layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace = False)]
                else:
                    layers += [conv2d, nn.ReLU(inplace = False)]

                in_channels = v

        return nn.ModuleList(layers)

def vgg11(pretrained = False, **kwargs):

    if pretrained:
        kwargs['init_weights'] = False
    model = VGG('A', **kwargs)

    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg11']))
    
    return model


def vgg11_bn(pretrained=False, **kwargs):
    """VGG 11-layer model (configuration "A") with batch normalization
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    if pretrained:
        kwargs['init_weights'] = False
    model = VGG('A', batch_norm=True, **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg11_bn']))
    return model


def vgg13(pretrained=False, **kwargs):
    """VGG 13-layer model (configuration "B")
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    if pretrained:
        kwargs['init_weights'] = False
    model = VGG('B', **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg13']))
    return model


def vgg13_bn(pretrained=False, **kwargs):
    """VGG 13-layer model (configuration "B") with batch normalization
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    if pretrained:
        kwargs['init_weights'] = False
    model = VGG('B', batch_norm=True, **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg13_bn']))
    return model


def vgg16(pretrained=False, **kwargs):
    """VGG 16-layer model (configuration "D")
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    if pretrained:
        kwargs['init_weights'] = False
    model = VGG('D', **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg16']))
    return model


def vgg16_bn(pretrained=False, **kwargs):
    """VGG 16-layer model (configuration "D") with batch normalization
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    if pretrained:
        kwargs['init_weights'] = False
    model = VGG('D', batch_norm=True, **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg16_bn']))
    return model


def vgg19(pretrained=False, **kwargs):
    """VGG 19-layer model (configuration "E")
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    if pretrained:
        kwargs['init_weights'] = False
    model = VGG('E', **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg19']))
    return model


def vgg19_bn(pretrained=False, **kwargs):
    """VGG 19-layer model (configuration 'E') with batch normalization
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    if pretrained:
        kwargs['init_weights'] = False
    model = VGG('E', batch_norm=True, **kwargs)
    if pretrained:
        model.load_state_dict(model_zoo.load_url(model_urls['vgg19_bn']))

    return model

if __name__ == "__main__":
    vgg = vgg11_bn(pretrained = True)
    test = torch.randn((1, 3, 224, 224))
    vgg(test)
    
