import torch
import torchvision
import torchvision.transforms as transforms

import matplotlib.pyplot as plt
import numpy as np

import torch.optim as optim

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

#import dla
import vgg



EPOCHS = 40

class CnnBlock(nn.Module):
    def __init__(self, in_channels=3, out_channels=16):
        super(CnnBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, 2*in_channels, 3, 1)
        self.conv2 = nn.Conv2d(2*in_channels, out_channels, 3, 1)
        self.bn1 = nn.BatchNorm2d(2*in_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = F.relu(x)
        return x

class CifarCNN(nn.Module):
    def __init__(self):
        super(CifarCNN, self).__init__()
        self.block1 = CnnBlock(3, 16)
        self.block2 = CnnBlock(16, 64)
        self.block3 = CnnBlock(64, 128)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.block1(x)
        x = self.dropout1(x)
        x = F.max_pool2d(x,2)
        x = self.block2(x)
        x = F.max_pool2d(x,2)
        x = self.dropout2(x)
        x = self.block3(x)
        x = torch.flatten(x,1)
        x = self.fc1(x)
        output = F.log_softmax(x, dim=1)
        return output

def train(model, trainloader, device, optimizer, criterion, epoch):
        running_loss = 0.0
        for i, [data, target] in enumerate(trainloader, 0):
            data, target = data.to(device), target.to(device)
            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            # print statistics
            running_loss += loss.item()
            if i % 2000 == 1999:    # print every 2000 mini-batches
                print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
                running_loss = 0.0


def test(model, testloader, device, classes):
    correct = 0
    total = 0

    # since we're not training, we don't need to calculate the gradients for our outputs
    with torch.no_grad():
        for data, target in testloader:
            data, target = data.to(device), target.to(device)
            # calculate outputs by running images through the network
            output = model(data)
            # the class with the highest energy is what we choose as prediction
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()

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


    # prepare to count predictions for each class
    correct_pred = {classname: 0 for classname in classes}
    total_pred = {classname: 0 for classname in classes}

    # again no gradients needed
    with torch.no_grad():
        for data, target in testloader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            _, predictions = torch.max(output, 1)
            # collect the correct predictions for each class
            for label, prediction in zip(target, predictions):
                if label == prediction:
                    correct_pred[classes[label]] += 1
                total_pred[classes[label]] += 1


    # print accuracy for each class
    for classname, correct_count in correct_pred.items():
        accuracy = 100 * float(correct_count) / total_pred[classname]
        print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')


def main():
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    batch_size = 4

    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                            download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                              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=batch_size,
                                             shuffle=False, num_workers=2)

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    model = vgg.VGG('VGG16').to(device)

    optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=0.9)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(EPOCHS):
        train(model, trainloader, device, optimizer, criterion, epoch)
        test(model, testloader, device, classes)

    PATH = './cifar_vgg.pth'
    torch.save(model.state_dict(), PATH)


if __name__ == "__main__":
    main()