import torch
import torch.nn as nn

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']
}


def conv_Ternarize(W):
    with torch.no_grad():
        m = W.abs().sum() / W.numel()
        W = torch.clamp(torch.round(W / m), min=-1, max=1)
        return W


class TnConv2d(nn.Conv2d):
    def __init__(self, *args, **kwargs):
        super(TnConv2d, self).__init__(*args, **kwargs)

    def forward(self, x):
        w = self.weight
        w_tn = w + (conv_Ternarize(w.data) - w).detach()
        output = nn.functional.conv2d(x, w_tn, self.bias, self.stride,
                                      self.padding, self.dilation, self.groups)
        return output


def linear_Ternarize(W):
    with torch.no_grad():
        m = W.abs().mean()
        W = torch.clamp(torch.round(W / m), min=-1, max=1)
        return W


class TnLinear(nn.Linear):
    def __init__(self, *args, **kwargs):
        super(TnLinear, self).__init__(*args, **kwargs)

    def forward(self, x):
        w = self.weight
        w_tn = w + (linear_Ternarize(w.data) - w).detach()
        output = nn.functional.linear(x, w_tn, self.bias)
        return output


class VGG(nn.Module):
    def __init__(self, features, ternary, num_class=100):
        super().__init__()
        self.features = features
        self.classifier = nn.Sequential(
            TnLinear(512, 4096) if ternary else nn.Linear(512, 4096),
            nn.BatchNorm1d(4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            TnLinear(4096, 4096) if ternary else nn.Linear(4096, 4096),
            nn.BatchNorm1d(4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            TnLinear(4096, num_class) if ternary else nn.Linear(4096, num_class),
            nn.BatchNorm1d(num_class),
        )

    def forward(self, x):
        output = self.features(x)
        output = output.view(output.size()[0], -1)
        output = self.classifier(output)

        return output


def make_layers(cfg, ternary, batch_norm=False):
    layers = []

    input_channel = 3
    for l in cfg:
        if l == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            continue

        layers += [TnConv2d(input_channel, l, kernel_size=3, padding=1)] if ternary else [nn.Conv2d(input_channel, l, kernel_size=3, padding=1)]

        if batch_norm:
            layers += [nn.BatchNorm2d(l)]

        layers += [nn.ReLU(inplace=True)]
        input_channel = l

    return nn.Sequential(*layers)


def vgg11(ternary=False):
    return VGG(make_layers(cfg['A'], ternary, batch_norm=False), ternary)


def vgg11_bn(ternary=False):
    return VGG(make_layers(cfg['A'], ternary, batch_norm=True), ternary)


def vgg13(ternary=False):
    return VGG(make_layers(cfg['B'], ternary, batch_norm=False), ternary)


def vgg13_bn(ternary=False):
    return VGG(make_layers(cfg['B'], ternary, batch_norm=True), ternary)


def vgg16(ternary=False):
    return VGG(make_layers(cfg['D'], ternary, batch_norm=False), ternary)


def vgg16_bn(ternary=False):
    return VGG(make_layers(cfg['D'], ternary, batch_norm=True), ternary)


def vgg19(ternary=False):
    return VGG(make_layers(cfg['E'], ternary, batch_norm=False), ternary)


def vgg19_bn(ternary=False):
    return VGG(make_layers(cfg['E'], ternary, batch_norm=True), ternary)
