'''
From https://github.com/fangchangma/sparse-to-dense.pytorch
Modified to work with our data augmentation method and dense noisy RGB-D images.
'''

import os
import torch
import shutil
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

cmap = plt.cm.hot

def parse_command():
	model_names = ['resnet18', 'resnet50']
	loss_names = ['l1', 'l2']
	data_names = ['gt_nyudepthv2', 'nyudepthv2', 'kitti']
	from dataloaders.dense_to_sparse import UniformSampling, SimulatedStereo, SimulatedCameraNoise, AlgorithmicNoise
	sparsifier_names = [x.name for x in [UniformSampling, SimulatedStereo, SimulatedCameraNoise, AlgorithmicNoise]] + ["None"]
	from models import Decoder
	decoder_names = Decoder.names
	from dataloaders.dataloader import MyDataloader
	modality_names = MyDataloader.modality_names

	import argparse
	parser = argparse.ArgumentParser(description='Sparse-to-Dense')
	parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18', choices=model_names,
						help='model architecture: ' + ' | '.join(model_names) + ' (default: resnet18)')
	parser.add_argument('--data', metavar='DATA', default='nyudepthv2',
						help='dataset: ' + ' | '.join(data_names) + ' (default: nyudepthv2)' +
						' Or full pathname to data directory.')
	parser.add_argument('--modality', '-m', metavar='MODALITY', default='rgbd', choices=modality_names,
						help='modality: ' + ' | '.join(modality_names) + ' (default: rgbd)')
	parser.add_argument('-s', '--num-samples', default=0, type=int, metavar='N',
						help='number of sparse depth samples (default: 0)')
	parser.add_argument('--max-depth', default=-1.0, type=float, metavar='D',
						help='cut-off depth of sparsifier, negative values means infinity (default: inf [m])')
	parser.add_argument('--sparsifier', metavar='SPARSIFIER', default=UniformSampling.name, choices=sparsifier_names,
						help='sparsifier: ' + ' | '.join(sparsifier_names) + ' (default: ' + SimulatedCameraNoise.name + ')')
	parser.add_argument('--sparsemodel', '-sm', default='', help='name of sparsifier model. (default: none)')
	parser.add_argument('--numaugmented', '-na', default=0, type=int, help='the number of augmented frames per scene (default: 0 for no augmented data)')
	parser.add_argument('--simoffline', action='store_true', help='use static augmented examples as if they had been pre-generated')
	parser.add_argument('--decoder', '-d', metavar='DECODER', default='deconv5', choices=decoder_names,
						help='decoder: ' + ' | '.join(decoder_names) + ' (default: deconv5)')
	parser.add_argument('-j', '--workers', default=10, type=int, metavar='N',
						help='number of data loading workers (default: 10)')
	parser.add_argument('--epochs', default=15, type=int, metavar='N',
						help='number of total epochs to run (default: 15)')
	parser.add_argument('-c', '--criterion', metavar='LOSS', default='l1', choices=loss_names,
						help='loss function: ' + ' | '.join(loss_names) + ' (default: l1)')
	parser.add_argument('-b', '--batch-size', default=8, type=int, help='mini-batch size (default: 8)')
	parser.add_argument('--lr', '--learning-rate', default=0.01, type=float,
						metavar='LR', help='initial learning rate (default 0.01)')
	parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
						help='momentum')
	parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
						metavar='W', help='weight decay (default: 1e-4)')
	parser.add_argument('--print-freq', '-p', 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', type=str, default='',
						help='evaluate model on validation set')
	parser.add_argument('--gpu', '-gpu', default=0, type=int,
						metavar='N', help='which gpu to use (default: 0)')
	parser.add_argument('--fname', '-f', help='name of output folder. If none, is automatically generated.')
	parser.add_argument('--no-pretrain', dest='pretrained', action='store_false',
						help='not to use ImageNet pre-trained weights')
	parser.add_argument('--samplecap', default=-1, type=int, help='the number of noisy depth frames to sample per scene (default: -1 for all)')
	parser.add_argument('--scenecap', default=-1, type=int, help='the number of scenes to sample (default: -1 for all)')
	parser.set_defaults(pretrained=True)
	args = parser.parse_args()
	if args.modality == 'rgb' and args.num_samples != 0:
		print("number of samples is forced to be 0 when input modality is rgb")
		args.num_samples = 0
	if args.modality == 'rgb' and args.max_depth != 0.0:
		print("max depth is forced to be 0.0 when input modality is rgb/rgbd")
		args.max_depth = 0.0
	return args

def save_checkpoint(state, is_best, epoch, output_directory):
	checkpoint_filename = os.path.join(output_directory, 'checkpoint-' + str(epoch) + '.pth.tar')
	torch.save(state, checkpoint_filename)
	if is_best:
		best_filename = os.path.join(output_directory, 'model_best.pth.tar')
		shutil.copyfile(checkpoint_filename, best_filename)
	if epoch > 0:
		prev_checkpoint_filename = os.path.join(output_directory, 'checkpoint-' + str(epoch-1) + '.pth.tar')
		if os.path.exists(prev_checkpoint_filename):
			os.remove(prev_checkpoint_filename)

def adjust_learning_rate(optimizer, epoch, lr_init):
	"""Sets the learning rate to the initial LR decayed by 10 every 5 epochs"""
	lr = lr_init * (0.1 ** (epoch // 5))
	for param_group in optimizer.param_groups:
		param_group['lr'] = lr

def get_output_directory(args):
	output_directory = os.path.join('results',
		'{}.sparsifier={}.samples={}.modality={}.arch={}.decoder={}.criterion={}.lr={}.bs={}.pretrained={}'.
		format(args.data, args.sparsifier, args.num_samples, args.modality, \
			args.arch, args.decoder, args.criterion, args.lr, args.batch_size, \
			args.pretrained))
	return output_directory


def colored_depthmap(depth, d_min=None, d_max=None):
	if d_min is None:
		d_min = np.min(depth)
	if d_max is None:
		d_max = np.max(depth)
	depth_relative = (depth - d_min) / (d_max - d_min)
	return 255 * cmap(depth_relative)[:,:,:3] # H, W, C


def merge_into_row(input, depth_target, depth_pred):
	rgb = 255 * np.transpose(np.squeeze(input.cpu().numpy()), (1,2,0)) # H, W, C
	depth_target_cpu = np.squeeze(depth_target.cpu().numpy())
	depth_pred_cpu = np.squeeze(depth_pred.data.cpu().numpy())

	d_min = min(np.min(depth_target_cpu), np.min(depth_pred_cpu))
	d_max = max(np.max(depth_target_cpu), np.max(depth_pred_cpu))
	depth_target_col = colored_depthmap(depth_target_cpu, d_min, d_max)
	depth_pred_col = colored_depthmap(depth_pred_cpu, d_min, d_max)
	img_merge = np.hstack([rgb, depth_target_col, depth_pred_col])

	return img_merge


def merge_into_row_with_gt(input, depth_input, depth_target, depth_pred):
	rgb = 255 * np.transpose(np.squeeze(input.cpu().numpy()), (1,2,0)) # H, W, C
	depth_input_cpu = np.squeeze(depth_input.cpu().numpy())
	depth_target_cpu = np.squeeze(depth_target.cpu().numpy())
	depth_pred_cpu = np.squeeze(depth_pred.data.cpu().numpy())

	d_min = min(np.min(depth_input_cpu), np.min(depth_target_cpu), np.min(depth_pred_cpu))
	d_max = max(np.max(depth_input_cpu), np.max(depth_target_cpu), np.max(depth_pred_cpu))
	depth_input_col = colored_depthmap(depth_input_cpu, d_min, d_max)
	depth_target_col = colored_depthmap(depth_target_cpu, d_min, d_max)
	depth_pred_col = colored_depthmap(depth_pred_cpu, d_min, d_max)

	img_merge = np.hstack([rgb, depth_input_col, depth_target_col, depth_pred_col])

	return img_merge


def add_row(img_merge, row):
	return np.vstack([img_merge, row])


def save_image(img_merge, filename):
	img_merge = Image.fromarray(img_merge.astype('uint8'))
	img_merge.save(filename)