'''
Modified from https://github.com/chengyangfu/pytorch-vgg-cifar10
'''
import math
from torch import batch_norm

import torch
import torch.nn as nn
import torch.nn.init as init
from collections import OrderedDict
from models.vgg import VGG
from models.nbn import NoisyBatchNorm2d

# __all__ = [
#     'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn',
#     'vgg19_bn', 'vgg19',, 'vgg'
# ]


class VGGNBN(VGG):
    '''
    VGG model 
    '''
    def __init__(self, features, num_classes=10):
        super(VGGNBN, self).__init__(features, num_classes)

def make_layers(cfg, norm_type=None):
    layers = []
    in_channels = 3
    mp_idx = 1 # maxpool index
    ft_idx = 1 # feature components

    for v in cfg:
        if v == 'M':
            layers += [('maxpool%d'%mp_idx, nn.MaxPool2d(kernel_size=2, stride=2))]
            mp_idx += 1
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            layers += [('conv%d'%ft_idx,conv2d), ('bn%d'%ft_idx,NoisyBatchNorm2d(v)), ('act%d'%ft_idx,nn.ReLU(inplace=True))]
            ft_idx += 1
            in_channels = v
    return nn.Sequential(OrderedDict(layers))


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'],
    '8': [64, 'M', 128, 'M', 256, 'M', 512,'M', 512, 'M'], # small
    'S': [32, 'M', 64, 'M', 128, 'M', 256,'M', 512, 'M'], # small
    'SS': [16, 'M', 32, 'M', 64, 'M', 128,'M', 512, 'M'], # small
    'SSS': [8, 'M', 16, 'M', 32, 'M', 64,'M', 512, 'M'], # small
}

def vggnbn11(num_classes=10):
    """VGG 11-layer model (configuration "A")"""
    return VGGNBN(make_layers(cfg['A']), num_classes)

def vggnbn13(num_classes=10):
    """VGG 13-layer model (configuration "A")"""
    return VGGNBN(make_layers(cfg['B']), num_classes)

def vggnbn16(num_classes=10):
    """VGG 16-layer model (configuration "A")"""
    return VGGNBN(make_layers(cfg['D']), num_classes)

def vggnbn19(num_classes=10):
    """VGG 19-layer model (configuration "E")"""
    return VGGNBN(make_layers(cfg['E']), num_classes)


if __name__ == "__main__":
    model = vggnbn16()
    inputs = torch.randn(128,3,32,32)
    outputs = model(inputs)
    print(outputs.shape)