r"""Main script to train models."""

from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

import os
import sys
import shutil
import numpy as np
from tqdm import tqdm

import torch
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable
from sklearn.metrics import mean_squared_error

from model import AdaptorHigh   # best FC6 model
from model import AdaptorMid    # best CONV42 model
from model import AdaptorEarly  # best POOL1 model
from dataset import VisualDataset
from globals import TRAIN_TEST_DIR


def save_checkpoint(state, is_best, folder='./', filename='checkpoint.pth.tar'):
    r"""Save a snapshot of the model.

    @param state: torch state_dict object
    @param is_best: boolean
                    whether this is our best performing model yet
    @param folder: string [default: ./]
                   which folder to save to
    @param filename: string [default: checkpoint.pth.tar]
                     name of checkpoint
    """
    if not os.path.isdir(folder):
        os.makedirs(folder)
    torch.save(state, os.path.join(folder, filename))
    if is_best:
        shutil.copyfile(os.path.join(folder, filename),
                        os.path.join(folder, 'model_best.pth.tar'))


def load_checkpoint(file_path, use_cuda=False):
    r"""Given a saved checkpoint (see save_checkpoint function), we initialize
    a model from the saved weights.

    @param file_path: string
                      path to saved model path
    @param use_cuda: boolean [default: False]
                     whether parameters are by default CUDA casted
    """
    checkpoint = torch.load(file_path) if use_cuda else \
        torch.load(file_path, map_location=lambda storage, location: storage)

    if checkpoint['adaptor'] == 'high':
        assert 'hiddens_dim' in checkpoint
        model = AdaptorHigh(hiddens_dim=checkpoint['hiddens_dim'])
    elif checkpoint['adaptor'] == 'mid':
        model = AdaptorMid()
    elif checkpoint['adaptor'] == 'early':
        model = AdaptorLow()
    else:
        raise Exception('Unrecognized adaptor: %s' % checkpoint['adaptor'])

    model.load_state_dict(checkpoint['state_dict'])
    model.adaptor = checkpoint['adaptor']

    return model


class AverageMeter(object):
    r"""Computes and stores the average and current value"""
    def __init__(self):
        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 cross_entropy(input, target, eps=1e-6):
    r"""k-Class Cross Entropy (Log Softmax + Log Loss)

    @param input: torch.Tensor (size N x K)
    @param target: torch.Tensor (size N x K)
    @param eps: error to add (default: 1e-6)
    @return loss: torch.Tensor (size N)
    """
    if not (target.size(0) == input.size(0)):
        raise ValueError(
            "Target size ({}) must be the same as input size ({})".format(
                target.size(0), input.size(0)))

    log_input = F.log_softmax(input + eps, dim=1)
    loss = target * log_input
    return -torch.mean(torch.sum(loss, dim=1))


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('adaptor', type=str, help='high|mid|low')
    # we pick 128 as the default b/c it seems the smallest dimensionality as which we can
    # still get good performance.
    parser.add_argument('--hiddens-dim', type=int, default=128, help='number of hidden dims [default: 128]')
    # this multiplier is helpful to keep gradients large but training stable.
    parser.add_argument('--loss-scale', type=float, default=10000., help='multiplier for loss [default: 10000.]')
    # by default, we use the first split
    parser.add_argument('--train-test-split-dir', type=str, default=os.path.join(TRAIN_TEST_DIR, '1'),
                        help='where to load train_test_split paths [default: ./train_test_split/1]')
    parser.add_argument('--out-dir', type=str, default='./trained_models',
                        help='where to save checkpoints [./trained_models]')
    parser.add_argument('--batch-size', type=int, default=10,
                        help='number of examples in a mini-batch [default: 10]')
    parser.add_argument('--lr', type=float, default=1e-4, help='learning rate [default: 1e-4]')
    parser.add_argument('--epochs', type=int, default=100, help='number of epochs [default: 100]')
    parser.add_argument('--cuda', action='store_true', default=False)
    args = parser.parse_args()
    args.cuda = args.cuda and torch.cuda.is_available()

    if not os.path.isdir(args.out_dir):
        os.makedirs(args.out_dir)

    train_dataset = VisualDataset(adaptor=args.adaptor, split='train',
                                  train_test_split_dir=args.train_test_split_dir)
    val_dataset = VisualDataset(adaptor=args.adaptor, split='val',
                                train_test_split_dir=args.train_test_split_dir)
    test_dataset = VisualDataset(adaptor=args.adaptor, split='test',
                                 train_test_split_dir=args.train_test_split_dir)

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)

    if args.adaptor == 'high':
        model = AdaptorHigh(args.hiddens_dim)
    elif args.adaptor == 'mid':
        model = AdaptorMid()
    elif args.adaptor == 'early':
        model = AdaptorEarly()
    else:
        raise Exception('Unrecognized adaptor: %s' % args.adaptor)

    if args.cuda:
        model.cuda()

    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4)

    def train(epoch):
        model.train()
        loss_meter = AverageMeter()

        for batch_idx, (sketch, label) in enumerate(train_loader):
            sketch = Variable(sketch)
            label = Variable(label.float())
            batch_size = len(sketch)

            if args.cuda:
                sketch = sketch.cuda()
                label = label.cuda()

            # set optimizer defaults to 0
            optimizer.zero_grad()

            pred_logits = []
            photo_generator = train_dataset.gen_photos()
            for photo in photo_generator():
                photo = Variable(photo)
                if args.cuda:
                    photo = photo.cuda()
                photo = (photo.repeat(batch_size, 1) if args.adaptor == 'high' else
                         photo.repeat(batch_size, 1, 1, 1))
                pred_logit = model(photo, sketch)
                pred_logits.append(pred_logit)

            pred_logits = torch.cat(pred_logits, dim=1)
            loss = args.loss_scale * cross_entropy(pred_logits, label)
            loss_meter.update(loss.data[0], batch_size)

            loss.backward()
            optimizer.step()
            mean_grads = torch.mean(torch.cat([param.grad.cpu().data.contiguous().view(-1)
                                               for param in model.parameters()]))
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\t|Grad|: {:6f}'.format(
                epoch, batch_idx * batch_size, len(train_loader.dataset),  100. * batch_idx / len(train_loader),
                loss_meter.avg, mean_grads))

        print('====> Epoch: {}\tLoss: {:.4f}'.format(epoch, loss_meter.avg))

        return loss_meter.avg

    def validate():
        model.eval()
        loss_meter = AverageMeter()
        pbar = tqdm(total=len(val_loader))

        for batch_idx, (sketch, label) in enumerate(val_loader):
            sketch = Variable(sketch, volatile=True)
            label = Variable(label.float(), requires_grad=False)
            batch_size = len(sketch)

            if args.cuda:
                sketch = sketch.cuda()
                label = label.cuda()

            pred_logits = []
            photo_generator = val_dataset.gen_photos()
            for photo in photo_generator():
                photo = Variable(photo)
                if args.cuda:
                    photo = photo.cuda()
                photo = (photo.repeat(batch_size, 1) if args.adaptor == 'high' else
                         photo.repeat(batch_size, 1, 1, 1))
                pred_logit = model(photo, sketch)
                pred_logits.append(pred_logit)

            pred_logits = torch.cat(pred_logits, dim=1)
            loss = args.loss_scale * cross_entropy(pred_logits, label)
            loss_meter.update(loss.data[0], batch_size)

            pbar.update()
        pbar.close()

        print('====> Val Loss: {:.4f}'.format(loss_meter.avg))
        return loss_meter.avg


    def test():
        model.eval()
        loss_meter = AverageMeter()
        pbar = tqdm(total=len(test_loader))

        for batch_idx, (sketch, label) in enumerate(test_loader):
            sketch = Variable(sketch, volatile=True)
            label = Variable(label.float(), requires_grad=False)
            batch_size = len(sketch)

            if args.cuda:
                sketch = sketch.cuda()
                label = label.cuda()

            pred_logits = []
            photo_generator = test_dataset.gen_photos()
            for photo in photo_generator():
                photo = Variable(photo)
                if args.cuda:
                    photo = photo.cuda()
                photo = (photo.repeat(batch_size, 1) if args.adaptor == 'high' else
                         photo.repeat(batch_size, 1, 1, 1))
                pred_logit = model(photo, sketch)
                pred_logits.append(pred_logit)

            pred_logits = torch.cat(pred_logits, dim=1)
            loss = args.loss_scale * cross_entropy(pred_logits, label)
            loss_meter.update(loss.data[0], batch_size)

            pbar.update()
        pbar.close()

        print('====> Test Loss: {:.4f}'.format(loss_meter.avg))
        return loss_meter.avg


    best_loss = sys.maxint
    store_loss = np.zeros((args.epochs, 3))

    for epoch in xrange(1, args.epochs + 1):
        train_loss = train(epoch)
        val_loss = validate()
        test_loss = test()

        is_best = val_loss < best_loss
        best_loss = min(val_loss, best_loss)

        save_checkpoint({
            'state_dict': model.state_dict(),
            'train_loss': train_loss,
            'val_loss': val_loss,
            'test_loss': test_loss,
            'adaptor': args.adaptor,
            'hiddens_dim': args.hiddens_dim,
            'optimizer' : optimizer.state_dict(),
        }, is_best, folder=args.out_dir)

        store_loss[epoch - 1, 0] = train_loss
        store_loss[epoch - 1, 1] = val_loss
        store_loss[epoch  - 1, 2] = test_loss

        np.save(os.path.join(args.out_dir, 'summary-loss.npy'), store_loss)

r"""Main script to train models."""

from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

import os
import sys
import shutil
import numpy as np
from tqdm import tqdm

import torch
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable
from sklearn.metrics import mean_squared_error

from model import AdaptorHigh   # best FC6 model
from model import AdaptorMid    # best CONV42 model
from model import AdaptorEarly  # best POOL1 model
from dataset import VisualDataset
from globals import TRAIN_TEST_DIR


def save_checkpoint(state, is_best, folder='./', filename='checkpoint.pth.tar'):
    r"""Save a snapshot of the model.

    @param state: torch state_dict object
    @param is_best: boolean
                    whether this is our best performing model yet
    @param folder: string [default: ./]
                   which folder to save to
    @param filename: string [default: checkpoint.pth.tar]
                     name of checkpoint
    """
    if not os.path.isdir(folder):
        os.makedirs(folder)
    torch.save(state, os.path.join(folder, filename))
    if is_best:
        shutil.copyfile(os.path.join(folder, filename),
                        os.path.join(folder, 'model_best.pth.tar'))


def load_checkpoint(file_path, use_cuda=False):
    r"""Given a saved checkpoint (see save_checkpoint function), we initialize
    a model from the saved weights.

    @param file_path: string
                      path to saved model path
    @param use_cuda: boolean [default: False]
                     whether parameters are by default CUDA casted
    """
    checkpoint = torch.load(file_path) if use_cuda else \
        torch.load(file_path, map_location=lambda storage, location: storage)

    if checkpoint['adaptor'] == 'high':
        assert 'hiddens_dim' in checkpoint
        model = AdaptorHigh(hiddens_dim=checkpoint['hiddens_dim'])
    elif checkpoint['adaptor'] == 'mid':
        model = AdaptorMid()
    elif checkpoint['adaptor'] == 'early':
        model = AdaptorLow()
    else:
        raise Exception('Unrecognized adaptor: %s' % checkpoint['adaptor'])

    model.load_state_dict(checkpoint['state_dict'])
    model.adaptor = checkpoint['adaptor']

    return model


class AverageMeter(object):
    r"""Computes and stores the average and current value"""
    def __init__(self):
        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 cross_entropy(input, target, eps=1e-6):
    r"""k-Class Cross Entropy (Log Softmax + Log Loss)

    @param input: torch.Tensor (size N x K)
    @param target: torch.Tensor (size N x K)
    @param eps: error to add (default: 1e-6)
    @return loss: torch.Tensor (size N)
    """
    if not (target.size(0) == input.size(0)):
        raise ValueError(
            "Target size ({}) must be the same as input size ({})".format(
                target.size(0), input.size(0)))

    log_input = F.log_softmax(input + eps, dim=1)
    loss = target * log_input
    return -torch.mean(torch.sum(loss, dim=1))


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('adaptor', type=str, help='high|mid|low')
    # we pick 128 as the default b/c it seems the smallest dimensionality as which we can
    # still get good performance.
    parser.add_argument('--hiddens-dim', type=int, default=128, help='number of hidden dims [default: 128]')
    # this multiplier is helpful to keep gradients large but training stable.
    parser.add_argument('--loss-scale', type=float, default=10000., help='multiplier for loss [default: 10000.]')
    # by default, we use the first split
    parser.add_argument('--train-test-split-dir', type=str, default=os.path.join(TRAIN_TEST_DIR, '1'),
                        help='where to load train_test_split paths [default: ./train_test_split/1]')
    parser.add_argument('--out-dir', type=str, default='./trained_models',
                        help='where to save checkpoints [./trained_models]')
    parser.add_argument('--batch-size', type=int, default=10,
                        help='number of examples in a mini-batch [default: 10]')
    parser.add_argument('--lr', type=float, default=1e-4, help='learning rate [default: 1e-4]')
    parser.add_argument('--epochs', type=int, default=100, help='number of epochs [default: 100]')
    parser.add_argument('--cuda', action='store_true', default=False)
    args = parser.parse_args()
    args.cuda = args.cuda and torch.cuda.is_available()

    if not os.path.isdir(args.out_dir):
        os.makedirs(args.out_dir)

    train_dataset = VisualDataset(adaptor=args.adaptor, split='train',
                                  train_test_split_dir=args.train_test_split_dir)
    val_dataset = VisualDataset(adaptor=args.adaptor, split='val',
                                train_test_split_dir=args.train_test_split_dir)
    test_dataset = VisualDataset(adaptor=args.adaptor, split='test',
                                 train_test_split_dir=args.train_test_split_dir)

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)

    if args.adaptor == 'high':
        model = AdaptorHigh(args.hiddens_dim)
    elif args.adaptor == 'mid':
        model = AdaptorMid()
    elif args.adaptor == 'early':
        model = AdaptorEarly()
    else:
        raise Exception('Unrecognized adaptor: %s' % args.adaptor)

    if args.cuda:
        model.cuda()

    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4)

    def train(epoch):
        model.train()
        loss_meter = AverageMeter()

        for batch_idx, (sketch, label) in enumerate(train_loader):
            sketch = Variable(sketch)
            label = Variable(label.float())
            batch_size = len(sketch)

            if args.cuda:
                sketch = sketch.cuda()
                label = label.cuda()

            # set optimizer defaults to 0
            optimizer.zero_grad()

            pred_logits = []
            photo_generator = train_dataset.gen_photos()
            for photo in photo_generator():
                photo = Variable(photo)
                if args.cuda:
                    photo = photo.cuda()
                photo = (photo.repeat(batch_size, 1) if args.adaptor == 'high' else
                         photo.repeat(batch_size, 1, 1, 1))
                pred_logit = model(photo, sketch)
                pred_logits.append(pred_logit)

            pred_logits = torch.cat(pred_logits, dim=1)
            loss = args.loss_scale * cross_entropy(pred_logits, label)
            loss_meter.update(loss.data[0], batch_size)

            loss.backward()
            optimizer.step()
            mean_grads = torch.mean(torch.cat([param.grad.cpu().data.contiguous().view(-1)
                                               for param in model.parameters()]))
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\t|Grad|: {:6f}'.format(
                epoch, batch_idx * batch_size, len(train_loader.dataset),  100. * batch_idx / len(train_loader),
                loss_meter.avg, mean_grads))

        print('====> Epoch: {}\tLoss: {:.4f}'.format(epoch, loss_meter.avg))

        return loss_meter.avg

    def validate():
        model.eval()
        loss_meter = AverageMeter()
        pbar = tqdm(total=len(val_loader))

        for batch_idx, (sketch, label) in enumerate(val_loader):
            sketch = Variable(sketch, volatile=True)
            label = Variable(label.float(), requires_grad=False)
            batch_size = len(sketch)

            if args.cuda:
                sketch = sketch.cuda()
                label = label.cuda()

            pred_logits = []
            photo_generator = val_dataset.gen_photos()
            for photo in photo_generator():
                photo = Variable(photo)
                if args.cuda:
                    photo = photo.cuda()
                photo = (photo.repeat(batch_size, 1) if args.adaptor == 'high' else
                         photo.repeat(batch_size, 1, 1, 1))
                pred_logit = model(photo, sketch)
                pred_logits.append(pred_logit)

            pred_logits = torch.cat(pred_logits, dim=1)
            loss = args.loss_scale * cross_entropy(pred_logits, label)
            loss_meter.update(loss.data[0], batch_size)

            pbar.update()
        pbar.close()

        print('====> Val Loss: {:.4f}'.format(loss_meter.avg))
        return loss_meter.avg


    def test():
        model.eval()
        loss_meter = AverageMeter()
        pbar = tqdm(total=len(test_loader))

        for batch_idx, (sketch, label) in enumerate(test_loader):
            sketch = Variable(sketch, volatile=True)
            label = Variable(label.float(), requires_grad=False)
            batch_size = len(sketch)

            if args.cuda:
                sketch = sketch.cuda()
                label = label.cuda()

            pred_logits = []
            photo_generator = test_dataset.gen_photos()
            for photo in photo_generator():
                photo = Variable(photo)
                if args.cuda:
                    photo = photo.cuda()
                photo = (photo.repeat(batch_size, 1) if args.adaptor == 'high' else
                         photo.repeat(batch_size, 1, 1, 1))
                pred_logit = model(photo, sketch)
                pred_logits.append(pred_logit)

            pred_logits = torch.cat(pred_logits, dim=1)
            loss = args.loss_scale * cross_entropy(pred_logits, label)
            loss_meter.update(loss.data[0], batch_size)

            pbar.update()
        pbar.close()

        print('====> Test Loss: {:.4f}'.format(loss_meter.avg))
        return loss_meter.avg


    best_loss = sys.maxint
    store_loss = np.zeros((args.epochs, 3))

    for epoch in xrange(1, args.epochs + 1):
        train_loss = train(epoch)
        val_loss = validate()
        test_loss = test()

        is_best = val_loss < best_loss
        best_loss = min(val_loss, best_loss)

        save_checkpoint({
            'state_dict': model.state_dict(),
            'train_loss': train_loss,
            'val_loss': val_loss,
            'test_loss': test_loss,
            'adaptor': args.adaptor,
            'hiddens_dim': args.hiddens_dim,
            'optimizer' : optimizer.state_dict(),
        }, is_best, folder=args.out_dir)

        store_loss[epoch - 1, 0] = train_loss
        store_loss[epoch - 1, 1] = val_loss
        store_loss[epoch  - 1, 2] = test_loss

        np.save(os.path.join(args.out_dir, 'summary-loss.npy'), store_loss)

Main script to train models. # best FC6 model # best CONV42 model # best POOL1 model Save a snapshot of the model.      @param state: torch state_dict object     @param is_best: boolean                     whether this is our best performing model yet     @param folder: string [default: ./]                    which folder to save to     @param filename: string [default: checkpoint.pth.tar]                      name of checkpoint Given a saved checkpoint (see save_checkpoint function), we initialize     a model from the saved weights.      @param file_path: string                       path to saved model path     @param use_cuda: boolean [default: False]                      whether parameters are by default CUDA casted Computes and stores the average and current value k-Class Cross Entropy (Log Softmax + Log Loss)      @param input: torch.Tensor (size N x K)     @param target: torch.Tensor (size N x K)     @param eps: error to add (default: 1e-6)     @return loss: torch.Tensor (size N) # we pick 128 as the default b/c it seems the smallest dimensionality as which we can # still get good performance. # this multiplier is helpful to keep gradients large but training stable. # by default, we use the first split # set optimizer defaults to 0