import torch

import os

import copy

from tqdm import tqdm

import matplotlib.pyplot as plt

def save_losses_graph(train_losses, valid_losses, experiment_dir):

    epochs = range(len(train_losses))

    plt.plot(epochs, train_losses, 'g', label='Training loss')

    plt.plot(epochs, valid_losses, 'b', label='Validation loss')

    plt.title('Training and Validation loss')

    plt.xlabel('Epochs')

    plt.ylabel('Loss')

    plt.legend(loc = 'upper right')

    plt.savefig(os.path.join(experiment_dir, 'train_valid_loss.png'))

    plt.clf()

def save_acc_graph(train_accs, valid_accs, experiment_dir):

    epochs = range(len(train_accs))

    plt.plot(epochs, train_accs, 'g', label='Training acc')

    plt.plot(epochs, valid_accs, 'b', label='Validation acc')

    plt.title('Training and Validation acc')

    plt.xlabel('Epochs')

    plt.ylabel('Acc')

    plt.legend(loc = 'upper right')

    plt.savefig(os.path.join(experiment_dir, 'train_valid_acc.png'))

    plt.clf()

def train(conf):

    device = conf['device']

    model = conf['model'].to(device)

    dataloaders = conf['dataloaders']

    criterion = conf['criterion']

    optimizer = conf['optimizer']

    scheduler = conf['scheduler']

    num_epochs = conf['num_epochs']

    experiment_dir = conf['experiment_dir']

    #valid_acc_history = []

    best_acc = 0.0

    epoch_bar = tqdm(range(num_epochs), desc='Epoch',unit='epoch')

    train_losses = []

    valid_losses = []

    train_accs = []

    valid_accs = []

    train_loss = -1.0

    valid_loss = -1.0

    train_acc = -1.0

    valid_acc = -1.0

    for epoch in epoch_bar:

        # Each epoch has a training and validation phase

        for phase in ['train', 'valid']:

            if phase == 'train':

                model.train()  # Set model to training mode

            else:

                model.eval()   # Set model to evaluate mode

            running_loss = 0.0

            running_corrects = 0

            # Iterate over data.

            batch_bar = tqdm(dataloaders[phase],

                             desc='Batch',

                             unit='batch',

                             leave=False)

            batch_losses = []

            for inputs, labels in batch_bar:

                inputs = inputs.to(device)

                labels = labels.to(device)

                # zero the parameter gradients

                optimizer.zero_grad()

                # forward

                # track history if only in train

                with torch.set_grad_enabled(phase == 'train'):

                    outputs = model(inputs)

                    loss = criterion(outputs, labels)

                    _, preds = torch.max(outputs, 1)

                    if phase == 'train':

                        loss.backward()

                        optimizer.step()

                # statistics

                running_corrects += torch.sum(preds == labels.data)

                batch_loss = loss.item()

                running_loss += batch_loss

                batch_losses.append(batch_loss)

                batch_bar.set_postfix(phase=phase, batch_loss=batch_loss)

            epoch_loss = running_loss / len(dataloaders[phase])

            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)

            if phase == 'train':

                train_loss = epoch_loss

                train_acc = epoch_acc.item()

                train_losses.append(train_loss)

                train_accs.append(train_acc)

            else:

                valid_loss = epoch_loss

                valid_acc = epoch_acc.item()

                valid_losses.append(valid_loss)

                valid_accs.append(valid_acc)

                save_losses_graph(train_losses, valid_losses, experiment_dir)

                save_acc_graph(train_accs, valid_accs, experiment_dir)

            if phase == 'valid' and epoch_acc > best_acc:

                best_acc = epoch_acc

                best_weights = copy.deepcopy(model.state_dict())

                weights_file = 'best_weights.pt'

                weights_path = os.path.join(experiment_dir, weights_file)

                torch.save(best_weights, weights_path)

        epoch_bar.set_postfix(tloss=train_loss,tacc=train_acc,

                              vloss=valid_loss, vacc=valid_acc)

        scheduler.step()

    print('Best valid Acc: {:4f}'.format(best_acc))

if __name__ == '__main__':

    from config import get_config

    # Get config from conf.yaml

    conf = get_config('./conf/training.yaml')

    # Train model

    train(conf)