import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules import Module
from torch.nn.parameter import Parameter
from torch.autograd import Variable

import torchvision
import torchvision.transforms as transforms


# The output of torchvision datasets are PILImage images of range [0, 1].
# We transform them to Tensors of normalized range [-1, 1]
transform=transforms.Compose([transforms.ToTensor(),
                              transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                             ])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                          shuffle=False, num_workers=2)
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

class BatchReNorm1d(Module):

    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, rmax=3.0, dmax=5.0):
        super(BatchReNorm1d, self).__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        self.affine = affine
        self.rmax = rmax
        self.dmax = dmax

        if self.affine:
            self.weight = Parameter(torch.Tensor(num_features))
            self.bias = Parameter(torch.Tensor(num_features))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)

        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))
        self.register_buffer('r', torch.ones(num_features))
        self.register_buffer('d', torch.zeros(num_features))

        self.reset_parameters()

    def reset_parameters(self):
        self.running_mean.zero_()
        self.running_var.fill_(1)
        self.r.fill_(1)
        self.d.zero_()
        if self.affine:
            self.weight.data.uniform_()
            self.bias.data.zero_()

    def _check_input_dim(self, input):
        if input.size(1) != self.running_mean.nelement():
            raise ValueError('got {}-feature tensor, expected {}'
                             .format(input.size(1), self.num_features))

    def forward(self, input):
        self._check_input_dim(input)
        n = input.size()[0]

        if self.training:
            mean = torch.mean(input, dim=0)

            sum = torch.sum((input - mean.expand_as(input))**2, dim=0)
            if sum == 0 and self.eps == 0:
                invstd = 0.0
            else:
                invstd = 1./torch.sqrt(sum/n + self.eps)
            unbiased_var = sum/(n - 1)


            self.r = torch.clamp(torch.sqrt(unbiased_var).data / torch.sqrt(self.running_var),
                            1./self.rmax, self.rmax)
            self.d = torch.clamp((mean.data - self.running_mean)/ torch.sqrt(self.running_var),
                            -self.dmax, self.dmax)

            r = Variable(self.r, requires_grad=False).expand_as(input)
            d = Variable(self.d, requires_grad=False).expand_as(input)

            input_normalized = (input - mean.expand_as(input)) * invstd.expand_as(input)

            input_normalized = input_normalized*r + d

            self.running_mean += self.momentum * (mean.data - self.running_mean)
            self.running_var += self.momentum * (unbiased_var.data - self.running_var)

            if not self.affine:
                return input_normalized

            output = input_normalized * self.weight.expand_as(input)
            output += self.bias.unsqueeze(0).expand_as(input)

            return output

        else:
            mean = Variable(self.running_mean).expand_as(input)
            invstd  = 1./ torch.sqrt(Variable(self.running_var).expand_as(input) + self.eps)

            input_normalized = (input - mean.expand_as(input)) * invstd.expand_as(input)

            if not self.affine:
                return input_normalized

            output = input_normalized * self.weight.expand_as(input)
            output += self.bias.unsqueeze(0).expand_as(input)

            return output

    def __repr__(self):
        return ('{name}({num_features}, eps={eps}, momentum={momentum},'
                'affine={affine}, rmax={rmax}, dmax={dmax})'
                .format(name=self.__class__.__name__, **self.__dict__))


class NetBatchReNorm(Module):
    def __init__(self):
        super(NetBatchReNorm, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.bn2d_1 = nn.BatchNorm2d(6)
        self.pool  = nn.MaxPool2d(2,2)

        self.conv2 = nn.Conv2d(6, 16, 5)
        self.bn2d_2   = nn.BatchNorm2d(16)

        self.fc1   = nn.Linear(16*5*5, 120)
        self.bn1d_1   = BatchReNorm1d(120)


        self.fc2   = nn.Linear(120, 84)
        self.bn1d_2 = BatchReNorm1d(84, rmax=1.0, dmax=0.0)

        self.fc3   = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(self.bn2d_1(F.relu(self.conv1(x))))
        x = self.pool(self.bn2d_2(F.relu(self.conv2(x))))
        x = x.view(-1, 16*5*5)
        x = self.bn1d_1(F.relu(self.fc1(x)))
        x = self.bn1d_2(F.relu(self.fc2(x)))
        x = self.fc3(x)
        return x

net2 = NetBatchReNorm()

import torch.optim as optim
criterion = nn.CrossEntropyLoss() # use a Classification Cross-Entropy loss
optimizer = optim.SGD(net2.parameters(), lr=0.001, momentum=0.9)

loss_no_batchnorm = []
for epoch in range(1): # loop over the dataset multiple times
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):

        # get the inputs
        inputs, labels = data
        # wrap them in Variable
        inputs, labels = Variable(inputs), Variable(labels)
        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net2(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.data[0]
        if i % 2000 == 1999: # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' % (epoch+1, i+1, running_loss / 2000))
            loss_no_batchnorm.append(running_loss / 2000)
            running_loss = 0.0
print('Finished Training')

net2.eval()

correct = 0
total = 0
for data in testloader:
    images, labels = data
    outputs = net2(Variable(images))
    _, predicted = torch.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted == labels).sum()

print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))