import os
import sys

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from matplotlib import pyplot as plt
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader

# progress bar
from tqdm.auto import tqdm

# hyper parameter
train_batch_size = 64
test_batch_size = 1000
epochs = 15
lr = 0.0003
gamma = 0.7
myseed = 1220


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = torch.flatten(x, end_dim=0)
        return output


def run():
    # CUDA, macOS GPU or CPU ?
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print('Using cuda')
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        print('Using mps')
    else:
        device = torch.device("cpu")
        print('Using cpu')

    # add random seed
    seed = myseed
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    # download data
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    dataset1 = datasets.MNIST('./data', train=True, download=True, transform=transform)
    dataset2 = datasets.MNIST('./data', train=False, transform=transform)
    train_loader = DataLoader(dataset1, batch_size=train_batch_size, shuffle=True, num_workers=1, pin_memory=True)
    test_loader = DataLoader(dataset2, batch_size=test_batch_size, shuffle=True, num_workers=1, pin_memory=True)

    model = Net().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)
    criterion = nn.CrossEntropyLoss()
    scheduler = StepLR(optimizer, step_size=1, gamma=gamma)

    # record the performance
    epoches = []
    training_accuracies = []
    training_loss = []
    testing_accuracies = []
    testing_loss = []
    training_accuracies.append('Training accuracies')
    training_loss.append('Training loss')
    testing_accuracies.append('Testing accuracies')
    testing_loss.append('Testing loss')

    # ---------- train and validate ----------
    for epoch in range(1, epochs + 1):
        total_train_loss, total_train_acc = 0.0, 0.0
        total_test_loss, total_test_acc = 0.0, 0.0
        train_batch_idx, test_batch_idx = 0, 0
        tqdm.write("[ epoch " + str(epoch) + " ]")
        epoches.append(epoch)

        # ---------- train ----------
        model.train()

        for batch in tqdm(train_loader, file=sys.stdout):

            train_batch_idx += 1
            if train_batch_idx == 1:
                tqdm.write("Training")

            imgs, labels = batch

            outputs = model(imgs.to(device))

            loss = criterion(outputs, labels.to(device))

            optimizer.zero_grad()

            loss.backward()

            nn.utils.clip_grad_norm_(model.parameters(), max_norm=10)

            optimizer.step()

            acc = (outputs.argmax(dim=-1) == labels.to(device)).float().mean()

            total_train_loss += loss.item()
            total_train_acc += acc.item()

        train_loss = total_train_loss / train_batch_idx
        train_acc = total_train_acc / train_batch_idx

        # ---------- validate ----------
        model.eval()

        for batch in tqdm(test_loader, file=sys.stdout):
            if test_batch_idx == 1:
                tqdm.write("Testing")

            test_batch_idx += 1

            imgs, labels = batch

            with torch.no_grad():
                outputs = model(imgs.to(device))

            loss = criterion(outputs, labels.to(device))

            acc = (outputs.argmax(dim=-1) == labels.to(device)).float().mean()

            total_test_loss += loss.item()
            total_test_acc += acc.item()

        test_loss = total_test_loss / test_batch_idx
        test_acc = total_test_acc / test_batch_idx

        training_accuracies.append(train_acc)
        training_loss.append(train_loss)
        testing_accuracies.append(test_acc)
        testing_loss.append(test_loss)
        print('Training Loss:', training_loss[epoch], 'Training Accuracy:', (100 * training_accuracies[epoch]),
              '%')
        print('Testing Loss:', testing_loss[epoch], 'Testing Accuracy:', 100 * testing_accuracies[epoch], '%')
        print('')

        scheduler.step()

    plot(epoches, training_loss)
    plot(epoches, training_accuracies)
    plot(epoches, testing_accuracies)
    plot(epoches, testing_loss)


def plot(epoches, performance):
    label = performance.pop(0)
    plt.title(label)
    plt.plot(epoches, performance, label=label)
    plt.xlabel('epoches')
    plt.legend()
    plt.savefig(label + '.jpg')
    plt.show()


if __name__ == '__main__':
    run()

