#!/usr/bin/env python3

#

# Note -- this training script is tweaked from the original at:

#           https://github.com/pytorch/examples/tree/master/imagenet

#

# For a step-by-step guide to transfer learning with PyTorch, see:

#           https://pytorch.org/tutorials/beginner/finetuning_torchvision_models_tutorial.html

#

import argparse

import os

import random

import time

import shutil

import warnings

import datetime

import torch

import torch.nn as nn

import torch.nn.parallel

import torch.nn.functional as F

import torch.backends.cudnn as cudnn

import torch.optim

import torch.utils.data

import torchvision.transforms as transforms

import torchvision.datasets as datasets

import torchvision.models as models

from torch.utils.tensorboard import SummaryWriter

from voc import VOCDataset

from nuswide import NUSWideDataset

from reshape import reshape_model

# get the available network architectures

model_names = sorted(name for name in models.__dict__

    if name.islower() and not name.startswith("__")

    and callable(models.__dict__[name]))

# parse command-line arguments

parser = argparse.ArgumentParser(description='PyTorch Image Classifier Training')

parser.add_argument('data', metavar='DIR',

                    help='path to dataset')

parser.add_argument('--dataset-type', type=str, default='folder',

                    choices=['folder', 'nuswide', 'voc'],

                    help='specify the dataset type (default: folder)')

parser.add_argument('--multi-label', action='store_true',

                    help='multi-label model (aka image tagging)')

parser.add_argument('--multi-label-threshold', type=float, default=0.5,

                    help='confidence threshold for counting a prediction as correct')

parser.add_argument('--model-dir', type=str, default='models', 

                    help='path to desired output directory for saving model '

                    'checkpoints (default: models/)')

parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet18',

                    choices=model_names,

                    help='model architecture: ' + ' | '.join(model_names) + ' (default: resnet18)')

parser.add_argument('--resolution', default=224, type=int, metavar='N',

                    help='input NxN image resolution of model (default: 224x224) '

                         'note than Inception models should use 299x299')

parser.add_argument('-j', '--workers', default=2, type=int, metavar='N',

                    help='number of data loading workers (default: 2)')

parser.add_argument('--epochs', default=35, type=int, metavar='N',

                    help='number of total epochs to run')

parser.add_argument('--start-epoch', default=0, type=int, metavar='N',

                    help='manual epoch number (useful on restarts)')

parser.add_argument('-b', '--batch-size', default=8, type=int, metavar='N',

                    help='mini-batch size (default: 8)')

parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,

                    metavar='LR', help='initial learning rate', dest='lr')

parser.add_argument('--momentum', default=0.9, type=float, metavar='M',

                    help='momentum')

parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,

                    metavar='W', help='weight decay (default: 1e-4)',

                    dest='weight_decay')

parser.add_argument('-p', '--print-freq', default=10, type=int,

                    metavar='N', help='print frequency (default: 10)')

parser.add_argument('--resume', default='', type=str, metavar='PATH',

                    help='path to latest checkpoint (default: none)')

parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',

                    help='evaluate model on validation set')

parser.add_argument('--pretrained', dest='pretrained', action='store_true', default=True,

                    help='use pre-trained model')

parser.add_argument('--seed', default=None, type=int,

                    help='seed for initializing training')

parser.add_argument('--gpu', default=0, type=int,

                    help='GPU ID to use (default: 0)')

args = parser.parse_args()

# open tensorboard logger (to model_dir/tensorboard)

tensorboard = SummaryWriter(log_dir=os.path.join(args.model_dir, "tensorboard", f"{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"))

print(f"To start tensorboard run:  tensorboard --log-dir={os.path.join(args.model_dir, 'tensorboard')}")

# variable for storing the best model accuracy so far

best_accuracy = 0

def main(args):

    """

    Load dataset, setup model, and train for N epochs

    """

    global best_accuracy

    if args.seed is not None:

        random.seed(args.seed)

        torch.manual_seed(args.seed)

        cudnn.deterministic = True

        warnings.warn('You have chosen to seed training. '

                      'This will turn on the CUDNN deterministic setting, '

                      'which can slow down your training considerably! '

                      'You may see unexpected behavior when restarting '

                      'from checkpoints.')

    if args.gpu is not None:

        print(f"=> using GPU {args.gpu} ({torch.cuda.get_device_name(args.gpu)})")

    # setup data transformations

    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],

                                     std=[0.229, 0.224, 0.225])

    train_transforms = transforms.Compose([

        transforms.RandomResizedCrop(args.resolution),

        transforms.RandomHorizontalFlip(),

        transforms.ToTensor(),

        normalize,

    ])

    val_transforms = transforms.Compose([

        transforms.Resize(args.resolution),

        transforms.CenterCrop(args.resolution),

        transforms.ToTensor(),

        normalize,

    ])

    # load the dataset

    if args.dataset_type == 'folder':

        train_dataset = datasets.ImageFolder(os.path.join(args.data, 'train'), train_transforms)

        val_dataset = datasets.ImageFolder(os.path.join(args.data, 'val'), val_transforms)

    elif args.dataset_type == 'nuswide':

        train_dataset = NUSWideDataset(args.data, 'trainval', train_transforms)

        val_dataset = NUSWideDataset(args.data, 'test', val_transforms)

    elif args.dataset_type == 'voc':

        train_dataset = VOCDataset(args.data, 'trainval', train_transforms)

        val_dataset = VOCDataset(args.data, 'val', val_transforms)

    if (args.dataset_type == 'nuswide' or args.dataset_type == 'voc') and (not args.multi_label):

        raise ValueError("nuswide or voc datasets should be run with --multi-label")

    print(f"=> dataset classes:  {len(train_dataset.classes)}  {train_dataset.classes}")

    train_loader = torch.utils.data.DataLoader(

        train_dataset, batch_size=args.batch_size, shuffle=True,

        num_workers=args.workers, pin_memory=True)

    val_loader = torch.utils.data.DataLoader(

        val_dataset, batch_size=args.batch_size, shuffle=False,

        num_workers=args.workers, pin_memory=True)

    # create or load the model if using pre-trained (the default)

    if args.pretrained:

        print(f"=> using pre-trained model '{args.arch}'")

        model = models.__dict__[args.arch](pretrained=True)

    else:

        print(f"=> creating model '{args.arch}'")

        model = models.__dict__[args.arch]()

    # reshape the model for the number of classes in the dataset

    model = reshape_model(model, args.arch, len(train_dataset.classes))

    # define loss function (criterion) and optimizer

    if args.multi_label:

        criterion = nn.BCEWithLogitsLoss()

    else:

        criterion = nn.CrossEntropyLoss()

    optimizer = torch.optim.SGD(model.parameters(), args.lr,

                                momentum=args.momentum,

                                weight_decay=args.weight_decay)

    # transfer the model to the GPU that it should be run on

    if args.gpu is not None:

        torch.cuda.set_device(args.gpu)

        model = model.cuda(args.gpu)

        criterion = criterion.cuda(args.gpu)

    # optionally resume from a checkpoint

    if args.resume:

        if os.path.isfile(args.resume):

            print(f"=> loading checkpoint '{args.resume}'")

            checkpoint = torch.load(args.resume)

            args.start_epoch = checkpoint['epoch'] + 1

            best_accuracy = checkpoint['best_accuracy']

            if args.gpu is not None:

                best_accuracy = best_accuracy.to(args.gpu)   # best_accuracy may be from a checkpoint from a different GPU

            model.load_state_dict(checkpoint['state_dict'])

            optimizer.load_state_dict(checkpoint['optimizer'])

            print(f"=> loaded checkpoint '{args.resume}' (epoch {checkpoint['epoch']})")

        else:

            print(f"=> no checkpoint found at '{args.resume}'")

    cudnn.benchmark = True

    # if in evaluation mode, only run validation

    if args.evaluate:

        validate(val_loader, model, criterion, 0)

        return

    # train for the specified number of epochs

    for epoch in range(args.start_epoch, args.epochs):

        # decay the learning rate

        adjust_learning_rate(optimizer, epoch)

        # train for one epoch

        train_loss, train_acc = train(train_loader, model, criterion, optimizer, epoch)

        # evaluate on validation set

        val_loss, val_acc = validate(val_loader, model, criterion, epoch)

        # remember best acc@1 and save checkpoint

        is_best = val_acc > best_accuracy

        best_accuracy = max(val_acc, best_accuracy)

        print(f"=> Epoch {epoch}")

        print(f"  * Train Loss     {train_loss:.4e}")

        print(f"  * Train Accuracy {train_acc:.4f}")

        print(f"  * Val Loss       {val_loss:.4e}")

        print(f"  * Val Accuracy   {val_acc:.4f}{'*' if is_best else ''}")

        save_checkpoint({

            'epoch': epoch,

            'arch': args.arch,

            'resolution': args.resolution,

            'classes': train_dataset.classes,

            'num_classes': len(train_dataset.classes),

            'multi_label': args.multi_label,

            'state_dict': model.state_dict(),

            'accuracy': {'train': train_acc, 'val': val_acc},

            'loss' : {'train': train_loss, 'val': val_loss},

            'optimizer' : optimizer.state_dict(),

        }, is_best)

def train(train_loader, model, criterion, optimizer, epoch):

    """

    Train one epoch over the dataset

    """

    batch_time = AverageMeter('Time', ':6.3f')

    data_time = AverageMeter('Data', ':6.3f')

    losses = AverageMeter('Loss', ':.4e')

    acc = AverageMeter('Accuracy', ':7.3f')

    progress = ProgressMeter(

        len(train_loader),

        [batch_time, data_time, losses, acc],

        prefix=f"Epoch: [{epoch}]")

    # switch to train mode

    model.train()

    # get the start time

    epoch_start = time.time()

    end = epoch_start

    # train over each image batch from the dataset

    for i, (images, target) in enumerate(train_loader):

        # measure data loading time

        data_time.update(time.time() - end)

        if args.gpu is not None:

            images = images.cuda(args.gpu, non_blocking=True)

            target = target.cuda(args.gpu, non_blocking=True)

        # compute output

        output = model(images)

        loss = criterion(output, target)

        # record loss and measure accuracy

        losses.update(loss.item(), images.size(0))

        acc.update(accuracy(output, target), images.size(0))

        # compute gradient and do SGD step

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

        # measure elapsed time

        batch_time.update(time.time() - end)

        end = time.time()

        if i % args.print_freq == 0 or i == len(train_loader)-1:

            progress.display(i)

    print(f"Epoch: [{epoch}] completed, elapsed time {time.time() - epoch_start:6.3f} seconds")

    tensorboard.add_scalar('Loss/train', losses.avg, epoch)

    tensorboard.add_scalar('Accuracy/train', acc.avg, epoch)

    return losses.avg, acc.avg

def validate(val_loader, model, criterion, epoch):

    """

    Measure model performance across the val dataset

    """

    batch_time = AverageMeter('Time', ':6.3f')

    losses = AverageMeter('Loss', ':.4e')

    acc = AverageMeter('Accuracy', ':7.3f')

    progress = ProgressMeter(

        len(val_loader),

        [batch_time, losses, acc],

        prefix='Val:   ')

    # switch to evaluate mode

    model.eval()

    with torch.no_grad():

        end = time.time()

        for i, (images, target) in enumerate(val_loader):

            if args.gpu is not None:

                images = images.cuda(args.gpu, non_blocking=True)

                target = target.cuda(args.gpu, non_blocking=True)

            # compute output

            output = model(images)

            loss = criterion(output, target)

            # record loss and measure accuracy

            losses.update(loss.item(), images.size(0))

            acc.update(accuracy(output, target), images.size(0))

            # measure elapsed time

            batch_time.update(time.time() - end)

            end = time.time()

            if i % args.print_freq == 0 or i == len(val_loader)-1:

                progress.display(i)

    tensorboard.add_scalar('Loss/val', losses.avg, epoch)

    tensorboard.add_scalar('Accuracy/val', acc.avg, epoch)

    return losses.avg, acc.avg

def save_checkpoint(state, is_best, filename='checkpoint.pth.tar', best_filename='model_best.pth.tar', labels_filename='labels.txt'):

    """

    Save a model checkpoint file, along with the best-performing model if applicable

    """

    if args.model_dir:

        model_dir = os.path.expanduser(args.model_dir)

        if not os.path.exists(model_dir):

            os.mkdir(model_dir)

        filename = os.path.join(model_dir, filename)

        best_filename = os.path.join(model_dir, best_filename)

        labels_filename = os.path.join(model_dir, labels_filename)

    # save the checkpoint

    torch.save(state, filename)

    # earmark the best checkpoint

    if is_best:

        shutil.copyfile(filename, best_filename)

        print(f"saved best model to:  {best_filename}")

    else:

        print(f"saved checkpoint to:  {filename}")

    # save labels.txt on the first epoch

    if state['epoch'] == 0:

        with open(labels_filename, 'w') as file:

            for label in state['classes']:

                file.write(f"{label}\n")

        print(f"saved class labels to:  {labels_filename}")

def adjust_learning_rate(optimizer, epoch):

    """

    Sets the learning rate to the initial LR decayed by 10 every 30 epochs

    """

    lr = args.lr * (0.1 ** (epoch // 30))

    for param_group in optimizer.param_groups:

        param_group['lr'] = lr

def accuracy(output, target):

    """

    Computes the accuracy of predictions vs groundtruth

    """

    with torch.no_grad():

        if args.multi_label:

            output = F.sigmoid(output)

            preds = ((output >= args.multi_label_threshold) == target.bool())   # https://medium.com/@yrodriguezmd/tackling-the-accuracy-multi-metric-9e2356f62513

            # https://stackoverflow.com/a/61585551

            #output[output >= args.multi_label_threshold] = 1

            #output[output < args.multi_label_threshold] = 0

            #preds = (output == target)

        else:

            output = F.softmax(output, dim=-1)

            _, preds = torch.max(output, dim=-1)

            preds = (preds == target)

        return preds.float().mean().cpu().item() * 100.0

class AverageMeter(object):

    """

    Computes and stores the average and current value

    """

    def __init__(self, name, fmt=':f'):

        self.name = name

        self.fmt = fmt

        self.reset()

    def reset(self):

        self.val = 0

        self.avg = 0

        self.sum = 0

        self.count = 0

    def update(self, val, n=1):

        self.val = val

        self.sum += val * n

        self.count += n

        self.avg = self.sum / self.count

    def __str__(self):

        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'

        return fmtstr.format(**self.__dict__)

class ProgressMeter(object):

    """

    Progress metering

    """

    def __init__(self, num_batches, meters, prefix=""):

        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)

        self.meters = meters

        self.prefix = prefix

    def display(self, batch):

        entries = [self.prefix + self.batch_fmtstr.format(batch)]

        entries += [str(meter) for meter in self.meters]

        print('  '.join(entries))

    def _get_batch_fmtstr(self, num_batches):

        num_digits = len(str(num_batches // 1))

        fmt = '{:' + str(num_digits) + 'd}'

        return '[' + fmt + '/' + fmt.format(num_batches) + ']'

if __name__ == '__main__':

    main(args)