import torch
import torch.optim as optim
import torch.nn.functional as F
import torchvision
import matplotlib.pyplot as plt

from network import *

# Name of the network to train
NETWORK = SimpleNet  # SimpleNet, DeepNet or ConvNet

NETWORK_NAME = NETWORK.__name__.lower()
RESTART_FROM_MODEL = False

# Hyperparameters
N_EPOCHS = 4
BATCH_SIZE_TRAIN = 64
BATCH_SIZE_TEST = 1000
LEARNING_RATE = 0.1
MOMENTUM = 0.5
LOG_INTERVAL = 10

torch.backends.cudnn.enabled = False
torch.manual_seed(1)

# Train dataset
train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('data/',
        train=True,
        download=True,
        transform=torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize((0.0,), (1.0,))
        ])),
    batch_size=BATCH_SIZE_TRAIN, shuffle=True)

# Test dataset
test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('data/',
        train=False,
        download=True,
        transform=torchvision.transforms.Compose([
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize((0.0,), (1.0,))
        ])),
        batch_size=BATCH_SIZE_TEST, shuffle=True)

# Network
network = NETWORK()

# Optimizer
optimizer = optim.SGD(network.parameters(), lr=LEARNING_RATE, momentum=MOMENTUM)

# Load an existing model
if RESTART_FROM_MODEL:
    try:
        network.load_state_dict(torch.load(f'data/model-{NETWORK_NAME}.pth'))
        optimizer.load_state_dict(torch.load(f'data/optimizer-{NETWORK_NAME}.pth'))
    except FileNotFoundError:
        print(f'No model found for {NETWORK_NAME}')

train_losses = []
train_counter = []
test_losses = []
test_counter = []

# Training
def train(epoch):
    network.train()
    correct = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = network(data)
        pred = output.data.max(1, keepdim=True)[1]
        correct += pred.eq(target.data.view_as(pred)).sum()
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % LOG_INTERVAL == 0:
            print('Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            train_losses.append(loss.item())
            train_counter.append(
                (batch_idx*64) + ((epoch-1)*len(train_loader.dataset)))
            torch.save(network.state_dict(), f'data/model-{NETWORK_NAME}.pth')
            torch.save(optimizer.state_dict(), f'data/optimizer-{NETWORK_NAME}.pth')
    print('\nTrain: Accuracy: {}/{} ({:.0f}%)'.format(
            correct, len(train_loader.dataset),
            100.0 * correct / len(train_loader.dataset)))

# Test
def test(epoch):
    network.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = network(data)
            test_loss += F.nll_loss(output, target, size_average=False).item()
            pred = output.data.max(1, keepdim=True)[1]
            correct += pred.eq(target.data.view_as(pred)).sum()
    test_loss /= len(test_loader.dataset)
    test_losses.append(test_loss)
    test_counter.append(len(train_loader.dataset) * epoch)
    print('Test: Accuracy: {}/{} ({:.0f}%)\n'.format(
            correct, len(test_loader.dataset),
            100.0 * correct / len(test_loader.dataset)))

# Main loop
test(0)
for epoch in range(1, N_EPOCHS + 1):
    train(epoch)
    test(epoch)

# Plot
fig = plt.figure()
plt.plot(train_counter, train_losses, color='blue', zorder=1)
plt.scatter(test_counter, test_losses, color='red', zorder=2)
plt.legend(['Train', 'Test'], loc='upper right')
plt.xlabel('Examples seen')
plt.ylabel('Loss')
fig.savefig(f'data/graph-{NETWORK_NAME}.png')



