"""
train.py
"""

import sys
import os
from collections import defaultdict
import numpy as np
import argparse
from tqdm import tqdm
from datetime import datetime

import torch
import torch.nn.functional as F
from torch import nn
from torch.autograd import Variable
import torchnet as tnt
from torchvision.utils import make_grid
from torch.optim import Adam
from torchnet.engine import Engine
from torchnet.logger import VisdomPlotLogger, VisdomLogger, VisdomTextLogger
import torchvision.models as models

from tensorboardX import SummaryWriter

parser = argparse.ArgumentParser()
# hyperparameters
parser.add_argument("--batch_size", default=100, type=int)
parser.add_argument("--num_routing_iterations", default=3, type=int)
parser.add_argument("--lr_init", default=0.001, type=float)
parser.add_argument("--lr_decay", default=0, type=float)
parser.add_argument("--momentum", default=0.9, type=float)
# affect operation of training/testing
parser.add_argument("--dataset", type=str, default="mnist", help="mnist, cifar10, fashion, svhn")
parser.add_argument("--transform", action="store_true", help="affinely transform data when testing")
parser.add_argument("--max_epochs", default=50, type=int)
parser.add_argument("-l", "--loading_epoch", type=int, help="Last saved parameters for resuming training")
parser.add_argument("--test", action="store_true")
# logging
parser.add_argument("--log_dir", default="logs")
parser.add_argument("--model_dir", default="epochs")
parser.add_argument("--tb_dir", default="tb")
parser.add_argument("--tag")
parser.add_argument("-t", "--track", action="store_true")
# system
parser.add_argument('--gpu', default=0, type=int)
parser.add_argument("--visdom_port", default=8097, type=int)
args = parser.parse_args()


# figure out names and if we're staring fresh (which means we'll rewrite logs and not load models)
name = "nr-"+ str(args.num_routing_iterations) + ("-trans" if args.transform else "") + ("-test" if args.test else "")
if args.tag is not None:
    name += "-" + args.tag
model_path = os.path.join(args.model_dir, args.dataset, name)
log_path = os.path.join(args.log_dir, args.dataset, name)
tb_path = os.path.join(args.tb_dir, args.dataset, name, datetime.now().isoformat().replace(":", "-"))
visdom_env = args.dataset + "-" + name
if args.track:
    assert args.visdom_port is not None
    writer = SummaryWriter(tb_path)
    writer.add_text("hyperparameters/batch_size", str(args.batch_size))
    writer.add_text("hyperparameters/lr_init", str(args.lr_init))
    writer.add_text("hyperparameters/lr_decay", str(args.lr_decay))
    writer.add_text("hyperparameters/momentum", str(args.momentum))
starting_fresh = args.loading_epoch == None

# setup dirs if not created
if args.model_dir != '':
    os.makedirs(model_path, exist_ok=True)
if args.log_dir != '' and starting_fresh:
    os.makedirs(log_path, exist_ok=True)
    f = open(log_path + '/train.txt','w')
    f.close()
    f = open(log_path + '/test.txt','w')
    f.close()

# print info about this run to visdom if it's a continuation
if not starting_fresh and args.track:
    text_logger = VisdomTextLogger(port=args.visdom_port)
    text_logger.log("dataset: " + "________________ " + str(args.dataset) + " " \
                    "batch_size: "  + "________________ " + str(args.batch_size) + " "\
                    "num_routing_iterations " + "________ " + str(args.num_routing_iterations) + " "\
                    "loading_epoch " + "_____________ " + str(args.loading_epoch))

# configure gpu and dataset
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
if args.dataset in ['mnist', 'fashion']:
    img_channels = 1
    img_width = 28
elif args.dataset in ['cifar10', 'svhn']:
    img_channels = 3
    img_width = 32
else:
    raise ValueError
def classes(dataset):
    if dataset in ['mnist', 'fashion', 'cifar10', 'svhn']:
        return 10
    else:
        raise ValueError
num_classes = classes(args.dataset)

### specify model and its loss
from model import CapsuleLayer, CapsuleNet, CapsuleLoss
model = CapsuleNet(img_channels, num_classes, args.num_routing_iterations, img_width)
model.cuda()
if not starting_fresh:
    print("Loading " + model_path)
    model.load_state_dict(torch.load(model_path +'/epoch_%d.pt' % args.loading_epoch))
capsule_loss = CapsuleLoss()
optimizer = Adam(model.parameters(), \
                 lr=args.lr_init, betas=(args.momentum, 0.999), eps=1e-8, weight_decay=args.lr_decay)

### wrappers for metrics (loss, accuracy, confusion)
meter_loss = tnt.meter.AverageValueMeter()
meter_accuracy = tnt.meter.ClassErrorMeter(accuracy=True)
confusion_meter = tnt.meter.ConfusionMeter(num_classes, normalized=True)
def reset_meters():
    meter_accuracy.reset()
    meter_loss.reset()
    confusion_meter.reset()
    
### show logs in visdom and log them if track
train_loss_logger = VisdomPlotLogger('line', env=visdom_env, opts={'title': 'Train Loss'}, port=args.visdom_port)
train_error_logger = VisdomPlotLogger('line', env=visdom_env, opts={'title': 'Train Accuracy'}, port=args.visdom_port)
test_loss_logger = VisdomPlotLogger('line', env=visdom_env, opts={'title': 'Test Loss'}, port=args.visdom_port)
test_accuracy_logger = VisdomPlotLogger('line', env=visdom_env, opts={'title': 'Test Accuracy'}, port=args.visdom_port)
confusion_logger = VisdomLogger('heatmap', env=visdom_env, opts={'title': 'Confusion matrix',
                                                 'columnnames': list(range(num_classes)),
                                                 'rownames': list(range(num_classes))})
ground_truth_logger = VisdomLogger('image', env=visdom_env, opts={'title': 'Ground Truth'}, port=args.visdom_port)
reconstruction_logger = VisdomLogger('image', env=visdom_env, opts={'title': 'Reconstruction'}, port=args.visdom_port)
perturbation_sample_logger = VisdomLogger('image', env=visdom_env, opts={'title': 'Perturbation'}, port=args.visdom_port)

all_reconstruction_logger = VisdomLogger('image', env=visdom_env, opts={'title': 'Figure 3: Reconstruction of Good/Bad Predictions'}, port=args.visdom_port)


def embedding(sample, all_mat, all_metadata, all_label_img):
    processed = process(sample[0])
    data = Variable(processed).cuda()
    _, _, vecs = model(data, save_vecs=True)
    all_metadata.append(sample[1])
    all_label_img.append(processed.cpu())
    for j, vec in enumerate(vecs):
        all_mat[j].append(vec.data.cpu().view(vec.size(0), -1))

def on_start(state):
    if args.loading_epoch is not None:
        state['epoch'] = args.loading_epoch

def on_sample(state):
    state['sample'].append(state['train'])

def on_start_epoch(state):
    reset_meters()
    if args.track:
        state['iterator'] = tqdm(state['iterator'])

def on_forward(state):
    meter_accuracy.add(state['output'].data, torch.LongTensor(state['sample'][1]))
    meter_loss.add(state['loss'].data[0])
    confusion_meter.add(state['output'].data, torch.LongTensor(state['sample'][1]))

def on_end_epoch(state):
    # train
    msg = '[%s] [Epoch %d] Training Loss: %.4f (Accuracy: %.2f%%)' % (
        visdom_env, state['epoch'], meter_loss.value()[0], meter_accuracy.value()[0])
    if args.log_dir != '':
        f = open(log_path + '/train.txt','a')
        f.write(msg + "\n")
        f.close()
    if args.track:
        print(msg)
        train_loss_logger.log(state['epoch'], meter_loss.value()[0])
        writer.add_scalar("train/loss", meter_loss.value()[0], state['epoch'])
        train_error_logger.log(state['epoch'], meter_accuracy.value()[0])
        writer.add_scalar("train/accuracy", meter_accuracy.value()[0], state['epoch'])
    reset_meters()
    # test
    engine.test(processor, get_iterator(args.dataset, False, args.batch_size, trans=args.transform))
    msg = '[%s] [Epoch %d] Testing Loss: %.4f (Accuracy: %.2f%%)' % (
        visdom_env, state['epoch'], meter_loss.value()[0], meter_accuracy.value()[0])
    if args.track:
        print(msg)
        test_loss_logger.log(state['epoch'], meter_loss.value()[0])
        writer.add_scalar("test/loss", meter_loss.value()[0], state['epoch'])
        test_accuracy_logger.log(state['epoch'], meter_accuracy.value()[0])
        writer.add_scalar("test/accuracy", meter_accuracy.value()[0], state['epoch'])
        confusion_logger.log(confusion_meter.value())
        # reconstructions
        reconstruction_iter = iter(get_iterator(args.dataset, False, trans=args.transform)) # False sets value of train mode
        all_mat = defaultdict(list)
        all_metadata = []
        all_label_img = []
        for i in range(3 if args.test else 100): # Accumulate more examples for embedding
            test_sample = next(reconstruction_iter)
            if i == 0:
                ground_truth = process(test_sample[0])
                _, reconstruction, perturbations = model(Variable(ground_truth).cuda(), perturb=state['epoch']%args.batch_size)
                reconstruction = reconstruction.cpu().view_as(ground_truth).data
                size = list(ground_truth.size())
                size[0] = 16 * 11
                perturbation = perturbations.cpu().view(size).data
            embedding(test_sample, all_mat, all_metadata, all_label_img)
        all_metadata = torch.cat(all_metadata)
        all_label_img = torch.cat(all_label_img)
        for j in range(args.num_routing_iterations):
            all_mat[j] = torch.cat(all_mat[j])
            writer.add_embedding(
                all_mat[j],
                metadata=all_metadata,
                label_img=all_label_img,
                global_step=state['epoch']*100+j+1,
                tag="Ep {} Iter {}".format(state['epoch'], j+1),
            )
        combined_mat = torch.cat([all_mat[j] for j in range(args.num_routing_iterations)])
        combined_metadata = torch.cat([torch.ones(all_mat[j].size(0))*(j+1) for j in range(args.num_routing_iterations)])
        combined_label_img = all_label_img.repeat(args.num_routing_iterations, 1, 1, 1)
        writer.add_embedding(
            combined_mat,
            metadata=combined_metadata,
            label_img=combined_label_img,
            global_step=state['epoch']*100,
            tag="Ep {} Comb".format(state['epoch']),
        )

        gt_image = make_grid(ground_truth, nrow=int(args.batch_size ** 0.5), normalize=True, range=(0, 1))
        writer.add_image("groundtruth", gt_image, state['epoch'])
        ground_truth_logger.log(gt_image.numpy())
        r_image = make_grid(reconstruction, nrow=int(args.batch_size ** 0.5), normalize=True, range=(0, 1))
        writer.add_image("reconstruction", r_image, state['epoch'])
        reconstruction_logger.log(r_image.numpy())
        p_image = make_grid(perturbation, nrow=11, normalize=True, range=(0, 1))
        writer.add_image("perturbation", p_image, state['epoch'])
        perturbation_sample_logger.log(p_image.numpy())

        # fig 3: all reconstructions across all target classes
        num_good, num_bad = 1, 1
        good_samples, bad_samples = [], []
        for sample in iter(get_iterator(args.dataset, False, 1)):
            img, true_lbl = sample
            true_lbl = int(true_lbl[0])

            if len(good_samples) == num_good and len(bad_samples) == num_bad:
                break

            ground_truth = process(img)

            classes, reconstruction, all_reconstructions = model(Variable(ground_truth).cuda(), all_reconstructions=True)

            # get prediction
            _, max_length_indices = classes.max(dim=1)
            pred_lbl = max_length_indices.data[0]
    
            # for the given image, get all reconstructions for each class 
            all_reconstructions = all_reconstructions.cpu().view(num_classes, img_channels, img_width, img_width).data # reconstructions: [torch.FloatTensor of size 10x1x28x28]
            candidate = (ground_truth, true_lbl, pred_lbl, all_reconstructions)
    
            # collect good samples
            if len(good_samples) < num_good and true_lbl == pred_lbl:
                good_samples.append(candidate)

            if len(bad_samples) < num_bad and true_lbl != pred_lbl:
                bad_samples.append(candidate)

        cats = []
        for sample in good_samples + bad_samples:
            ground_truth, true_lbl, pred_lbl, all_reconstructions = sample
            cat = torch.cat([ground_truth, all_reconstructions])
            maximum = cat.max()
            if cat.size(1) == 1:
                cat = cat.repeat(1, 3, 1, 1)
            cat[true_lbl+1, 1, -2:] = maximum
            cat[true_lbl+1, 1, :2] = maximum
            cat[true_lbl+1, 1, :, -2:] = maximum
            cat[true_lbl+1, 1, :, :2] = maximum
            if pred_lbl != true_lbl:
                cat[pred_lbl+1, 0, -2:] = maximum
                cat[pred_lbl+1, 0, :2] = maximum
                cat[pred_lbl+1, 0, :, -2:] = maximum
                cat[pred_lbl+1, 0, :, :2] = maximum
            cats.append(cat)
        image = make_grid(torch.cat(cats), nrow=1+num_classes, normalize=True, range=(0,1))
        writer.add_image("reconstruction_all", image, state['epoch'])
        all_reconstruction_logger.log(image.numpy())

    if args.log_dir != '':
        f = open(log_path + '/test.txt','a')
        f.write(msg + "\n")
        f.close()

    torch.save(model.state_dict(), model_path +'/epoch_%d.pt' % state['epoch'])


### engine ties everything together

engine = Engine()
engine.hooks['on_start'] = on_start
engine.hooks['on_sample'] = on_sample
engine.hooks['on_start_epoch'] = on_start_epoch
engine.hooks['on_forward'] = on_forward
engine.hooks['on_end_epoch'] = on_end_epoch

from utils import augmentation, get_iterator

def process(data):
    if args.dataset in ['mnist', 'fashion']:
        data = data.unsqueeze(1)
    elif args.dataset == 'cifar10':
        data = data.permute(0, 3, 1, 2)
    elif args.dataset == 'svhn':
        pass # explicit
    else:
        raise ValueError
    assert torch.max(data) > 2 # To ensure everything needs to be scaled down
    return data.float() / 255.0

def processor(sample):
    data, labels, training = sample
    data = augmentation(process(data))
    labels = torch.LongTensor(labels)
    labels = torch.sparse.torch.eye(num_classes).index_select(dim=0, index=labels)
    data = Variable(data).cuda()
    labels = Variable(labels).cuda()
    if training:
        classes, reconstruction = model(data, labels)
    else:
        classes, reconstruction = model(data)
    loss = capsule_loss(data, labels, classes, reconstruction)
    return loss, classes

engine.train(processor, get_iterator(args.dataset, True, test=args.test), \
             maxepoch=args.max_epochs, optimizer=optimizer)
