# Source code for deeprobust.image.netmodels.densenet

    """
    This is an implementation of DenseNet model.

    Reference
    ---------
    ..[1]Huang, Gao, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q. Weinberger. "Densely connected convolutional networks." In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700-4708. 2017.
    ..[2]Original implementation: https://github.com/kuangliu/pytorch-cifar
    """
    import math

    import torch
    import torch.nn as nn
    import torch.nn.functional as F

    class Bottleneck(nn.Module):
        def __init__(self, in_planes, growth_rate):
            super(Bottleneck, self).__init__()
            self.bn1 = nn.BatchNorm2d(in_planes)
            self.conv1 = nn.Conv2d(in_planes, 4*growth_rate, kernel_size=1, bias=False)
            self.bn2 = nn.BatchNorm2d(4*growth_rate)
            self.conv2 = nn.Conv2d(4*growth_rate, growth_rate, kernel_size=3, padding=1, bias=False)

        def forward(self, x):
            out = self.conv1(F.relu(self.bn1(x)))
            out = self.conv2(F.relu(self.bn2(out)))
            out = torch.cat([out,x], 1)
            return out

    class Transition(nn.Module):
        def __init__(self, in_planes, out_planes):
            super(Transition, self).__init__()
            self.bn = nn.BatchNorm2d(in_planes)
            self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=1, bias=False)

        def forward(self, x):
            out = self.conv(F.relu(self.bn(x)))
            out = F.avg_pool2d(out, 2)
            return out

    class DenseNet(nn.Module):
        """DenseNet.

        """

        def __init__(self, block, nblocks, growth_rate=12, reduction=0.5, num_classes=10):
            super(DenseNet, self).__init__()
            self.growth_rate = growth_rate

            num_planes = 2*growth_rate
            self.conv1 = nn.Conv2d(3, num_planes, kernel_size=3, padding=1, bias=False)

            self.dense1 = self._make_dense_layers(block, num_planes, nblocks[0])
            num_planes += nblocks[0]*growth_rate
            out_planes = int(math.floor(num_planes*reduction))
            self.trans1 = Transition(num_planes, out_planes)
            num_planes = out_planes

            self.dense2 = self._make_dense_layers(block, num_planes, nblocks[1])
            num_planes += nblocks[1]*growth_rate
            out_planes = int(math.floor(num_planes*reduction))
            self.trans2 = Transition(num_planes, out_planes)
            num_planes = out_planes

            self.dense3 = self._make_dense_layers(block, num_planes, nblocks[2])
            num_planes += nblocks[2]*growth_rate
            out_planes = int(math.floor(num_planes*reduction))
            self.trans3 = Transition(num_planes, out_planes)
            num_planes = out_planes

            self.dense4 = self._make_dense_layers(block, num_planes, nblocks[3])
            num_planes += nblocks[3]*growth_rate

            self.bn = nn.BatchNorm2d(num_planes)
            self.linear = nn.Linear(num_planes, num_classes)

        def _make_dense_layers(self, block, in_planes, nblock):
            layers = []
            for i in range(nblock):
                layers.append(block(in_planes, self.growth_rate))
                in_planes += self.growth_rate
            return nn.Sequential(*layers)

        def forward(self, x):
            out = self.conv1(x)
            out = self.trans1(self.dense1(out))
            out = self.trans2(self.dense2(out))
            out = self.trans3(self.dense3(out))
            out = self.dense4(out)
            out = F.avg_pool2d(F.relu(self.bn(out)), 4)
            out = out.view(out.size(0), -1)
            out = self.linear(out)
            return out

    def DenseNet121():
        """DenseNet121.
        """
        return DenseNet(Bottleneck, [6,12,24,16], growth_rate=32)

    def DenseNet169():
        """DenseNet169.
        """
        return DenseNet(Bottleneck, [6,12,32,32], growth_rate=32)

    def DenseNet201():
        """DenseNet201.
        """
        return DenseNet(Bottleneck, [6,12,48,32], growth_rate=32)

    def DenseNet161():
        """DenseNet161.
        """
        return DenseNet(Bottleneck, [6,12,36,24], growth_rate=48)

    def densenet_cifar():
        """densenet_cifar.
        """
        return DenseNet(Bottleneck, [6,12,24,16], growth_rate=12)

    def test(model, device, test_loader):
        """test.

        Parameters
        ----------
        model :
            model
        device :
            device
        test_loader :
            test_loader
        """
        model.eval()

        test_loss = 0
        correct = 0
        with torch.no_grad():
            for data, target in test_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                #test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss

                test_loss += F.cross_entropy(output, target)
                pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
                correct += pred.eq(target.view_as(pred)).sum().item()

        test_loss /= len(test_loader.dataset)

        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(test_loader.dataset),
            100. * correct / len(test_loader.dataset)))

    def train(model, device, train_loader, optimizer, epoch):
        """train.

        Parameters
        ----------
        model :
            model
        device :
            device
        train_loader :
            train_loader
        optimizer :
            optimizer
        epoch :
            epoch
        """
        model.train()

        # lr = util.adjust_learning_rate(optimizer, epoch, args) # don't need it if we use Adam

        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = torch.tensor(data).to(device), torch.tensor(target).to(device)
            optimizer.zero_grad()
            output = model(data)
            # loss = F.nll_loss(output, target)
            loss = F.cross_entropy(output, target)
            loss.backward()
            optimizer.step()
            if batch_idx % 10 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                    100. * batch_idx / len(train_loader), loss.item()/data.shape[0]))