import argparse
import numpy as np
import os
import sys

import shutil

import torch
import torch.optim as optim

from torch.utils.data import DataLoader

from tqdm import tqdm

import warnings

from lib.exceptions import NoGradientError
from lib.losses.lossPhotoTourism import loss_function
from lib.model import D2Net
from lib.dataloaders.datasetPhotoTourism_ipr import PhotoTourismIPR


# CUDA
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")

# Seed
torch.manual_seed(1)
if use_cuda:
	torch.cuda.manual_seed(1)
np.random.seed(1)

# Argument parsing
parser = argparse.ArgumentParser(description='Training script')

parser.add_argument(
	'--dataset_path', type=str, default="/scratch/udit/phototourism/",
	help='path to the dataset'
)

parser.add_argument(
	'--preprocessing', type=str, default='caffe',
	help='image preprocessing (caffe or torch)'
)

parser.add_argument(
	'--init_model', type=str, default='models/d2net.pth',
	help='path to the initial model'
)

parser.add_argument(
	'--num_epochs', type=int, default=10,
	help='number of training epochs'
)
parser.add_argument(
	'--lr', type=float, default=1e-3,
	help='initial learning rate'
)
parser.add_argument(
	'--batch_size', type=int, default=1,
	help='batch size'
)
parser.add_argument(
	'--num_workers', type=int, default=16,
	help='number of workers for data loading'
)

parser.add_argument(
	'--log_interval', type=int, default=250,
	help='loss logging interval'
)

parser.add_argument(
	'--log_file', type=str, default='log.txt',
	help='loss logging file'
)

parser.add_argument(
	'--plot', dest='plot', action='store_true',
	help='plot training pairs'
)
parser.set_defaults(plot=False)

parser.add_argument(
	'--checkpoint_directory', type=str, default='checkpoints',
	help='directory for training checkpoints'
)
parser.add_argument(
	'--checkpoint_prefix', type=str, default='rord',
	help='prefix for training checkpoints'
)

args = parser.parse_args()
print(args)

# Creating CNN model
model = D2Net(
	model_file=args.init_model,
	use_cuda=False
)
model = model.to(device)

# Optimizer
optimizer = optim.Adam(
	filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr
)

training_dataset = PhotoTourismIPR(
	base_path=args.dataset_path,
	preprocessing=args.preprocessing
)
training_dataset.build_dataset()

training_dataloader = DataLoader(
	training_dataset,
	batch_size=args.batch_size,
	num_workers=args.num_workers
)

# Define epoch function
def process_epoch(
		epoch_idx,
		model, loss_function, optimizer, dataloader, device,
		log_file, args, train=True, plot_path=None
):
	epoch_losses = []

	torch.set_grad_enabled(train)

	progress_bar = tqdm(enumerate(dataloader), total=len(dataloader))
	for batch_idx, batch in progress_bar:
		if train:
			optimizer.zero_grad()

		batch['train'] = train
		batch['epoch_idx'] = epoch_idx
		batch['batch_idx'] = batch_idx
		batch['batch_size'] = args.batch_size
		batch['preprocessing'] = args.preprocessing
		batch['log_interval'] = args.log_interval

		try:
			loss = loss_function(model, batch, device, plot=args.plot, plot_path=plot_path)
		except NoGradientError:
			# print("failed")
			continue

		current_loss = loss.data.cpu().numpy()[0]
		epoch_losses.append(current_loss)

		progress_bar.set_postfix(loss=('%.4f' % np.mean(epoch_losses)))

		if batch_idx % args.log_interval == 0:
			log_file.write('[%s] epoch %d - batch %d / %d - avg_loss: %f\n' % (
				'train' if train else 'valid',
				epoch_idx, batch_idx, len(dataloader), np.mean(epoch_losses)
			))

		if train:
			loss.backward()
			optimizer.step()

	log_file.write('[%s] epoch %d - avg_loss: %f\n' % (
		'train' if train else 'valid',
		epoch_idx,
		np.mean(epoch_losses)
	))
	log_file.flush()

	return np.mean(epoch_losses)


# Create the checkpoint directory
checkpoint_directory = os.path.join(args.checkpoint_directory, args.checkpoint_prefix)
if os.path.isdir(checkpoint_directory):
	print('[Warning] Checkpoint directory already exists.')
else:
	os.makedirs(checkpoint_directory, exist_ok=True)

# Open the log file for writing
log_file = os.path.join(checkpoint_directory,args.log_file)
if os.path.exists(log_file):
	print('[Warning] Log file already exists.')
log_file = open(log_file, 'a+')

# Create the folders for plotting if need be
plot_path=None
if args.plot:
	plot_path = os.path.join(checkpoint_directory,'train_vis')
	if os.path.isdir(plot_path):
		print('[Warning] Plotting directory already exists.')
	else:
		os.makedirs(plot_path, exist_ok=True)


# Initialize the history
train_loss_history = []

# Start the training
for epoch_idx in range(1, args.num_epochs + 1):
	# Process epoch
	train_loss_history.append(
		process_epoch(
			epoch_idx,
			model, loss_function, optimizer, training_dataloader, device,
			log_file, args, train=True, plot_path=plot_path
		)
	)

	# Save the current checkpoint
	checkpoint_path = os.path.join(
		checkpoint_directory,
		'%02d.pth' % (epoch_idx)
	)
	checkpoint = {
		'args': args,
		'epoch_idx': epoch_idx,
		'model': model.state_dict(),
		'optimizer': optimizer.state_dict(),
		'train_loss_history': train_loss_history,
	}
	torch.save(checkpoint, checkpoint_path)

# Close the log file
log_file.close()