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

import torch
import torch.nn as nn
import torch.nn.init as init
from models.ew import EWConv2d
from collections import OrderedDict

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


class EWVGG(nn.Module):
    '''
    VGG model 
    '''
    def __init__(self, features, num_classes=10):
        super(EWVGG, self).__init__()
        self.features = features
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(512, 512),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(512, 512),
            nn.ReLU(True),
            nn.Linear(512, num_classes),
        )
         # Initialize weights
        for m in self.modules():
            if isinstance(m, EWConv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                m.bias.data.zero_()


    def forward(self, x, mask=None):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

def get_norm(norm_type=None, v=0):
    if norm_type == None:
        return nn.Identity()
    elif norm_type == 'bn':
        return nn.BatchNorm2d(v)
    elif norm_type == 'in':
        return nn.GroupNorm(v,v)
    elif norm_type == 'ln':
        return nn.GroupNorm(1, v)
    elif norm_type.startswith('gn'):
        if len(norm_type)==2:
            num_groups = 32
        else:
            num_groups = int(norm_type[2:])
        return nn.GroupNorm(num_groups, v)
    else:
        raise NotImplementedError("Haven't implemented norm_type %s"%norm_type)

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 = EWConv2d(in_channels, v, kernel_size=3, padding=1)
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            layers += [('conv%d'%ft_idx,conv2d), ('bn%d'%ft_idx,get_norm(norm_type, 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 ewvgg8(num_classes=10):
    """VGG 8-layer model (configuration "8")"""
    return EWVGG(make_layers(cfg['8']), num_classes)

def ewvggs8(num_classes=10):
    """VGG 8-layer model (configuration "8")"""
    return EWVGG(make_layers(cfg['S']), num_classes)

def ewvggss8(num_classes=10):
    """VGG 8-layer model (configuration "8")"""
    return EWVGG(make_layers(cfg['SS']), num_classes)

def ewvggsss8(num_classes=10):
    """VGG 8-layer model (configuration "8")"""
    return EWVGG(make_layers(cfg['SSS']), num_classes)


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


def ewvgg11_bn():
    """VGG 11-layer model (configuration "A") with batch normalization"""
    return EWVGG(make_layers(cfg['A'], norm_type='bn'))


def ewvgg13():
    """VGG 13-layer model (configuration "B")"""
    return EWVGG(make_layers(cfg['B']))


def ewvgg13_bn():
    """VGG 13-layer model (configuration "B") with batch normalization"""
    return EWVGG(make_layers(cfg['B'], norm_type='bn'))


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


def ewvgg16_bn(num_classes=10):
    """VGG 16-layer model (configuration "D") with batch normalization"""
    return EWVGG(make_layers(cfg['D'], norm_type='bn'),num_classes)

# def vgg16_in(num_classes=10):
#     """VGG 16-layer model (configuration "D") with batch normalization"""
#     return VGG(make_layers(cfg['D'], norm_type='in'),num_classes)

# def vgg16_ln(num_classes=10):
#     """VGG 16-layer model (configuration "D") with batch normalization"""
#     return VGG(make_layers(cfg['D'], norm_type='ln'),num_classes)

# def vgg16_gn(num_classes=10):
#     """VGG 16-layer model (configuration "D") with batch normalization"""
#     return VGG(make_layers(cfg['D'], norm_type='gn'),num_classes)


# def vgg19():
#     """VGG 19-layer model (configuration "E")"""
#     return VGG(make_layers(cfg['E']))


# def vgg19_bn():
#     """VGG 19-layer model (configuration 'E') with batch normalization"""
#     return VGG(make_layers(cfg['E'], norm_type='bn'))

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