import os

from dataloader.datasetDHF1K import DHF1K
from torch.utils.data import DataLoader
from utils.salgan_utils import save_model, get_lr_optimizer
from utils.sendTelegram import send
from utils.printer import param_print
from utils.salgan_generator import create_model, add_bn
from evaluation.fast_evaluation import compute_metrics

import numpy as np

import torch
from torch.nn import AvgPool2d
from torch.nn.modules.loss import BCELoss
import torch.backends.cudnn as cudnn
from torch.optim import SGD, Adam
from torch.optim.lr_scheduler import ReduceLROnPlateau, StepLR
from time import time

from IPython import embed
from tensorboard_logger import configure, log_value, log_histogram


TRAIN = 'train'
VAL = 'val'
TEST = 'test'

def add_layer_weights(vgg_weights):
	# Mean of RGB weights of first layer with size [64,1,3,3]
	layer1 = vgg_weights['0.weight']
	mean_rgb = layer1.mean(dim=1,keepdim=True)
	vgg_weights['0.weight'] = torch.cat([layer1.cuda(),mean_rgb.cuda()],1)
	# We could do it easily accessing to the weights trought model[0].weight and change dimension 1, but as we
	# already have the 4th channel we'd be doing the mean of all of the channels, inicializing it in the wrong way.
	return vgg_weights

def train_eval(mode, model, optimizer, dataloader):
	if mode == TRAIN:
		N = len(ds_train)/batch_size
		model.train()
	else:
		N = len(ds_validate)/batch_size
		model.eval()

	total_loss = []
	#iterate epoch...
	#iterate epoch...
	for i, X in enumerate(dataloader[mode]):
		inputs = X[0].cuda()
		# noramlize saliency maps values between [0,1]
		gt_maps = X[1].cuda()/255
		embed()
		predictions = model.forward(inputs).squeeze()
		# reduce size for loss
		reduce_size = AvgPool2d((4,4))
		pred_ = reduce_size(predictions)
		gt_maps_ = reduce_size(gt_maps)

		pred_ = pred_.view(pred_.size()[0], -1)
		gt_maps_ = gt_maps_.view(gt_maps_.size()[0], -1)

		loss = bce_loss(pred_, gt_maps_)

		# make actual step update
		if mode==TRAIN:
			# compute gradients
			loss.backward()
			# step optimizer
			optimizer.step()
			# reset grads for next step
			optimizer.zero_grad()


		print("\t{}/{} loss:{}".format(i, int(N), loss.item()), end="\r")
		total_loss.append(loss.item())

	total_loss=np.mean(total_loss)
	return total_loss



if __name__ == '__main__':
	import argparse
	parser = argparse.ArgumentParser()
	parser.add_argument("--path_out", default='sal_dhf1k_adamdepthcoordaugm2_frombestsaldepth',
				type=str,
				help="""set output path for the trained model""")
	parser.add_argument("--batch_size", default=12,
				type=int,
				help="""Set batch size""")
	parser.add_argument("--n_epochs", default=10, type=int,
				help="""Set total number of epochs""")
	parser.add_argument("--depth", default=False, type=bool,
				help="""Enable 4th channel with depth""")
	parser.add_argument("--augment", default=False, type=bool,
				help="""Enable data augmentation""")
	parser.add_argument("--coord", default=False, type=bool,
				help="""Enable coordconv""")
	parser.add_argument("--flow", default=False, type=bool,
				help="""Enable opticalflow""")
	parser.add_argument("--lr", type=float, default=0.00001,
				help="""Learning rate for training""")
	parser.add_argument("--patience", type=int, default=3,
				help="""Patience for learning rate scheduler (default 10)""")
	args = parser.parse_args()


	# set output path ==========================================================
	path_out = '../trained_models/batch12_/' + args.path_out

	if not os.path.exists(path_out):
		# create output path
		os.makedirs(path_out)

		# create output for models
		path_models = os.path.join(path_out, 'models')
		if not os.path.exists(path_models):
			os.makedirs(path_models)

	# tensorboard
	configure("{}".format(path_out), flush_secs=5)

	# data =====================================================================
	batch_size = args.batch_size
	n_epochs = args.n_epochs
	lr = args.lr
	DEPTH = args.depth
	AUGMENT = args.augment
	COORD = args.coord
	FLOW = args.flow
	# Datasets for DHF1K
	ds_train = DHF1K(mode=TRAIN, transformation=True, depth=DEPTH, d_augm=AUGMENT, coord=COORD)
	ds_validate = DHF1K(mode=VAL, transformation=False, depth=DEPTH, d_augm=False, coord=COORD)

	# Dataloaders
	dataloader = {
		TRAIN: DataLoader(ds_train, batch_size=batch_size,
								shuffle=True, num_workers=2),
		VAL: DataLoader(ds_validate, batch_size=batch_size,
								shuffle=False, num_workers=2)
	}


	# POSSIBILITY OF CHOOSING GPU
	torch.cuda.set_device(1)
	# MODEL INITIALIZATION
	print("Init model...")
	vgg_weights = torch.load('../trained_models/salgan_baseline.pt')['state_dict']
	model = create_model(3)
	# if DEPTH and COORD:
	# 	model = create_model(6)
	# 	for i in range(0,3):
	# 		vgg_weights = add_layer_weights(vgg_weights)
	# elif DEPTH:
	# 	model = create_model(4)
	# 	add_layer_weights(vgg_weights)
	# elif COORD:
	# 	model = create_model(5)
	# 	for i in range(0,2):
	# 		vgg_weights = add_layer_weights(vgg_weights)
	# else: model = create_model(3)
	# Instead of adding manually the layer of new weights, we could use strict=False
	model.load_state_dict(vgg_weights)

	# Add batch normalization to current model if needed
	model = add_bn(model)

	model.train()
	model.cuda()
	cudnn.benchmark = True

	# NOT WORKING UNMOUNTED DISK
	# If we have the two GPU's available we are going to use both
	# if torch.cuda.device_count() > 1:
	# 	print("Using ", torch.cuda.device_count(), "GPUs!")
	# 	model = torch.nn.DataParallel(model)


	# LOSS FUNCTION
	bce_loss = BCELoss()

	# FINE-TUNE WHOLE NETWORK OR JUST DECODER => uncomment / or different lr for each part
	# decoder_parameters = []
	# base_params = []
	# for i, (a, p) in enumerate(model.named_parameters()):
	# 	embed()
	# 	if i>25:
	# 		# print(i, a, p.shape)
	# 		decoder_parameters.append(p)
	# 	else:
	# 		base_params.append(p)
			# If you wanna train just the decoder put this
			# p.requires_grad = False

	# ADAM OPTIMIZER
	optimizer = Adam(model.parameters(),
					lr = lr,
					weight_decay=0.000001)


	# STOCHASTIC GRADIENT DESCENT OPTIMIZER
	# optimizer = SGD(model.parameters(),
	# 				lr = 0.00001,
	# 				momentum=0.9,
	# 				weight_decay=0.00001,
	# 				nesterov=True)

	# NUMBER OF TOTAL PARAMETERS
	# pytorch_total_params = sum(p.numel() for p in model.parameters())
	# NUMBER OF TRAINABLE PARAMETERS
	trainable_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
	print("Trainable parameters: ", trainable_parameters)
	send("Trainable parameters: " + str(trainable_parameters))
	send("Experiment: " + args.path_out)

	# PRINT TABLE OF PARAMETERS
	param_print([path_out,"",DEPTH,AUGMENT,COORD,FLOW,batch_size,lr,n_epochs, trainable_parameters])

	# set learning rate scheduler
	# ReduceLROnPlateau(
		# optimizer,
		# mode (str) 'min':lr es reduira quan la metrica no es redueixi mes,  'max' al contrari,
		# factor (float) factor de reduccio de la lr,
		# patience (int) num epochs sense millora a partir dels quals es redueix lr,
		# verbose (bool),
	# )
	# scheduler = ReduceLROnPlateau(optimizer,
	# 							'min',
	# 							patience=args.patience,
	# 							verbose=True)
	scheduler = StepLR(optimizer, step_size=3, gamma=0.1)
	best_loss=9999999

	# main loop training =======================================================
	for id_epoch in range(n_epochs):
		for mode in [VAL, TRAIN]:
			# select dataloader
			data_iterator = dataloader[mode]
			#
			# # saliency metrics
			# if mode ==VAL:
			# 	print("Evaluating metrics....")
			# 	# only do 100 images from validation
			# 	metrics = compute_metrics(model, 100, DEPTH, COORD)
			#
			# 	# log metric values
			# 	for metric in metrics.keys():
			# 		log_value("Metrics/{}".format(metric),
			# 					metrics[metric], id_epoch)
			#
			# # get epoch loss
			# print("--> {} epoch {}".format(mode, id_epoch))


			epoch_loss = train_eval(mode, model, optimizer, dataloader)

			lr = list(get_lr_optimizer(optimizer))[0]
			print("-----------")
			print("Done! {} epoch {} loss {} lr {}".format(mode, id_epoch, epoch_loss, lr))
			send("{} epoch {}/{} loss {}".format(mode, id_epoch, n_epochs, epoch_loss))
			print("\n")

			# record loss
			log_value("loss/{}".format(mode), epoch_loss, id_epoch)
			log_value("lr/{}".format(mode), lr, id_epoch)
			# for v in model.state_dict():
			# 	log_histogram("Layer {}".format(v), model.state_dict()[v], id_epoch)
			if (id_epoch%2)==0:
				save_model(model, optimizer, id_epoch, path_out, name_model='{:03d}'.format(id_epoch))
			# store model if val loss improves
			if mode==VAL:
				if best_loss > epoch_loss:
					# update loss
					best_loss = epoch_loss

					save_model(model, optimizer, id_epoch, path_out, name_model='best')
				# scheduler.step(epoch_loss)
				scheduler.step()

[((2215, 2234), 'numpy.mean', 'np.mean', (['total_loss'], {}), '(total_loss)\n', (2222, 2234), True, 'import numpy as np\n'), ((2311, 2336), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (2334, 2336), False, 'import argparse\n'), ((4025, 4110), 'dataloader.datasetDHF1K.DHF1K', 'DHF1K', ([], {'mode': 'TRAIN', 'transformation': '(True)', 'depth': 'DEPTH', 'd_augm': 'AUGMENT', 'coord': 'COORD'}), '(mode=TRAIN, transformation=True, depth=DEPTH, d_augm=AUGMENT, coord=COORD\n    )\n', (4030, 4110), False, 'from dataloader.datasetDHF1K import DHF1K\n'), ((4121, 4198), 'dataloader.datasetDHF1K.DHF1K', 'DHF1K', ([], {'mode': 'VAL', 'transformation': '(False)', 'depth': 'DEPTH', 'd_augm': '(False)', 'coord': 'COORD'}), '(mode=VAL, transformation=False, depth=DEPTH, d_augm=False, coord=COORD)\n', (4126, 4198), False, 'from dataloader.datasetDHF1K import DHF1K\n'), ((4451, 4475), 'torch.cuda.set_device', 'torch.cuda.set_device', (['(1)'], {}), '(1)\n', (4472, 4475), False, 'import torch\n'), ((4613, 4628), 'utils.salgan_generator.create_model', 'create_model', (['(3)'], {}), '(3)\n', (4625, 4628), False, 'from utils.salgan_generator import create_model, add_bn\n'), ((5166, 5179), 'utils.salgan_generator.add_bn', 'add_bn', (['model'], {}), '(model)\n', (5172, 5179), False, 'from utils.salgan_generator import create_model, add_bn\n'), ((5493, 5502), 'torch.nn.modules.loss.BCELoss', 'BCELoss', ([], {}), '()\n', (5500, 5502), False, 'from torch.nn.modules.loss import BCELoss\n'), ((6509, 6545), 'utils.sendTelegram.send', 'send', (["('Experiment: ' + args.path_out)"], {}), "('Experiment: ' + args.path_out)\n", (6513, 6545), False, 'from utils.sendTelegram import send\n'), ((6577, 6685), 'utils.printer.param_print', 'param_print', (["[path_out, '', DEPTH, AUGMENT, COORD, FLOW, batch_size, lr, n_epochs,\n    trainable_parameters]"], {}), "([path_out, '', DEPTH, AUGMENT, COORD, FLOW, batch_size, lr,\n    n_epochs, trainable_parameters])\n", (6588, 6685), False, 'from utils.printer import param_print\n'), ((7118, 7159), 'torch.optim.lr_scheduler.StepLR', 'StepLR', (['optimizer'], {'step_size': '(3)', 'gamma': '(0.1)'}), '(optimizer, step_size=3, gamma=0.1)\n', (7124, 7159), False, 'from torch.optim.lr_scheduler import ReduceLROnPlateau, StepLR\n'), ((1601, 1608), 'IPython.embed', 'embed', ([], {}), '()\n', (1606, 1608), False, 'from IPython import embed\n'), ((1698, 1715), 'torch.nn.AvgPool2d', 'AvgPool2d', (['(4, 4)'], {}), '((4, 4))\n', (1707, 1715), False, 'from torch.nn import AvgPool2d\n'), ((3478, 3502), 'os.path.exists', 'os.path.exists', (['path_out'], {}), '(path_out)\n', (3492, 3502), False, 'import os\n'), ((3529, 3550), 'os.makedirs', 'os.makedirs', (['path_out'], {}), '(path_out)\n', (3540, 3550), False, 'import os\n'), ((3597, 3629), 'os.path.join', 'os.path.join', (['path_out', '"""models"""'], {}), "(path_out, 'models')\n", (3609, 3629), False, 'import os\n'), ((4240, 4312), 'torch.utils.data.DataLoader', 'DataLoader', (['ds_train'], {'batch_size': 'batch_size', 'shuffle': '(True)', 'num_workers': '(2)'}), '(ds_train, batch_size=batch_size, shuffle=True, num_workers=2)\n', (4250, 4312), False, 'from torch.utils.data import DataLoader\n'), ((4329, 4405), 'torch.utils.data.DataLoader', 'DataLoader', (['ds_validate'], {'batch_size': 'batch_size', 'shuffle': '(False)', 'num_workers': '(2)'}), '(ds_validate, batch_size=batch_size, shuffle=False, num_workers=2)\n', (4339, 4405), False, 'from torch.utils.data import DataLoader\n'), ((4539, 4589), 'torch.load', 'torch.load', (['"""../trained_models/salgan_baseline.pt"""'], {}), "('../trained_models/salgan_baseline.pt')\n", (4549, 4589), False, 'import torch\n'), ((3639, 3666), 'os.path.exists', 'os.path.exists', (['path_models'], {}), '(path_models)\n', (3653, 3666), False, 'import os\n'), ((3671, 3695), 'os.makedirs', 'os.makedirs', (['path_models'], {}), '(path_models)\n', (3682, 3695), False, 'import os\n'), ((7876, 7903), 'utils.salgan_utils.get_lr_optimizer', 'get_lr_optimizer', (['optimizer'], {}), '(optimizer)\n', (7892, 7903), False, 'from utils.salgan_utils import save_model, get_lr_optimizer\n'), ((8597, 8664), 'utils.salgan_utils.save_model', 'save_model', (['model', 'optimizer', 'id_epoch', 'path_out'], {'name_model': '"""best"""'}), "(model, optimizer, id_epoch, path_out, name_model='best')\n", (8607, 8664), False, 'from utils.salgan_utils import save_model, get_lr_optimizer\n')]