"""
Plot training results for regression and classification tasks on both 
sequence-mapped and residue-mapped data.

.............................................................................
idptools-parrot was developed by the Holehouse lab
     Original release ---- 2020

Question/comments/concerns? Raise an issue on github:
https://github.com/idptools/parrot

Licensed under the MIT license. 
"""

import numpy as np
import torch
import itertools
from scipy.stats import linregress
import matplotlib.pyplot as plt
import seaborn as sn
import pandas as pd

from parrot import encode_sequence

def training_loss(train_loss, val_loss, output_dir=''):
	"""Plot training and validation loss per epoch

	Figure is not displayed, but saved to file in current directory with the name
	'train_test.png'.

	Parameters
	----------
	train_loss : list
		training loss across each epoch
	val_loss : list
		validation loss across each epoch
	output_dir : str, optional
		directory to which the plot will be saved (default is current directory)
	"""

	fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(6,6))
	props = dict(boxstyle='round', facecolor='gainsboro', alpha=0.5)

	num_epochs = len(train_loss)

	# Loss per epoch
	training_loss, = ax.plot(np.arange(1, num_epochs+1), train_loss, label='Train')
	validation_loss, = ax.plot(np.arange(1, num_epochs+1), val_loss, label='Val')
	ax.set_xlabel("Epoch")
	ax.set_ylabel("Avg loss")
	ax.set_title("Training and testing loss per epoch")
	ax.legend(handles=[training_loss, validation_loss], fontsize=14, 
					facecolor='gainsboro', edgecolor='slategray')


	if num_epochs < 21:
		ax.set_xticks(np.arange(2, num_epochs+1, 2))
	elif num_epochs < 66:
		ax.set_xticks(np.arange(5, num_epochs+1, 5))
	elif num_epochs < 151:
		ax.set_xticks(np.arange(10, num_epochs+1, 10))
	else:
		ax.set_xticks(np.arange(50, num_epochs+1, 50))

	plt.savefig(output_dir + 'train_test.png')
	plt.clf()


def sequence_regression_scatterplot(true, predicted, output_dir=''):
	"""Create a scatterplot for a sequence-mapped values regression problem

	Figure is displayed to console if possible and saved to file in current 
	directory with the name 'seq_scatter.png'.

	Parameters
	----------
	true : list of PyTorch FloatTensors
		A list where each item is a [1 x 1] tensor with the true regression value
		of a particular sequence
	predicted : list of PyTorch FloatTensors
		A list where each item is a [1 x 1] tensor with the regression prediction
		for a particular sequence
	output_dir : str, optional
		directory to which the plot will be saved (default is current directory)
	"""

	true_list = []
	pred_list = []

	for item in true:
		true_list.append(item.cpu().numpy()[0][0])
	for item in predicted:
		pred_list.append(item.cpu().numpy()[0][0])

	plt.scatter(true_list, pred_list)
	edge_vals = [0.9*min(min(true_list), min(pred_list)), 
				 1.1*max(max(true_list), max(pred_list))]
	plt.xlim(edge_vals)
	plt.ylim(edge_vals)
	plt.plot(edge_vals, edge_vals, 'k--')
	plt.xlabel('True')
	plt.ylabel('Predicted')
	slope, intercept, r_value, p_value, std_err = linregress(true_list, pred_list)
	plt.title('Testing accuracy: R^2=%.3f' % (r_value**2))
	plt.savefig(output_dir + 'seq_scatter.png')
	plt.show()


def residue_regression_scatterplot(true, predicted, output_dir=''):
	"""Create a scatterplot for a residue-mapped values regression problem

	Each sequence is plotted with a unique marker-color combination, up to 70
	different sequences.

	Figure is displayed to console if possible and saved to file in current 
	directory with the name 'res_scatter.png'.

	Parameters
	----------
	true : list of PyTorch FloatTensors
		A list where each item is a [1 x len(sequence)] tensor with the true
		regression values of each residue in a sequence
	predicted : list of PyTorch FloatTensors
		A list where each item is a [1 x len(sequence)] tensor with the 
		regression predictions for each residue in a sequence
	output_dir : str, optional
		directory to which the plot will be saved (default is current directory)
	"""

	true_list = []
	pred_list = []

	marker = itertools.cycle(('>', '+', '.', 'o', '*', 'v', 'D')) 

	for item in true:
		single_frag = item.cpu().numpy()[0].flatten()
		true_list.append(list(single_frag))
	for item in predicted:
		single_frag = item.cpu().numpy()[0].flatten()
		pred_list.append(list(single_frag))

	for i in range(len(true_list)):
		plt.scatter(true_list[i], pred_list[i], s=6, marker=next(marker))

	plt.figure(1)

	left, right = plt.xlim()
	bottom, top = plt.ylim()
	edge_vals = [min(left, bottom), max(right, top)]
	plt.xlim(edge_vals)
	plt.ylim(edge_vals)
	plt.plot(edge_vals, edge_vals, 'k--')
	plt.xlabel('True')
	plt.ylabel('Predicted')
	slope, intercept, r_value, p_value, std_err = linregress(sum(true_list, []), sum(pred_list, []))
	plt.title('Testing accuracy: R^2=%.3f' % (r_value**2))
	plt.savefig(output_dir + 'res_scatter.png')
	plt.show()


def confusion_matrix(true_classes, predicted_classes, num_classes, output_dir=''):
	"""Create a confusion matrix for a sequence classification problem

	Figure is displayed to console if possible and saved to file in current 
	directory with the name 'seq_CM.png'.

	Parameters
	----------
	true_classes : list of PyTorch IntTensors
		A list where each item is a [1 x 1] tensor with the true class label of a
		particular sequence
	predicted_classes : list of PyTorch FloatTensors
		A list where each item is a [1 x num_classes] tensor prediction of the
		class label for a particular sequence
	num_classes : int
		Number of distinct data classes
	output_dir : str, optional
		directory to which the plot will be saved (default is current directory)
	"""

	cm = np.zeros((num_classes, num_classes))
	for i in range(len(true_classes)):
		cm[true_classes[i][0], np.argmax(predicted_classes[i][0].cpu().numpy())] += 1

	df_cm = pd.DataFrame(cm, range(num_classes), range(num_classes))
	sn.set(font_scale=1.4) # for label size
	sn.heatmap(df_cm, cmap='Blues', annot=True, annot_kws={"size": 16}) # font size
	plt.ylabel('True labels')
	plt.xlabel('Predicted labels')
	plt.title('Test set confusion matrix')
	plt.tight_layout()
	plt.savefig(output_dir + 'seq_CM.png')
	plt.show()


def res_confusion_matrix(true_classes, predicted_classes, num_classes, output_dir=''):
	"""Create a confusion matrix for a residue classification problem

	Figure is displayed to console if possible and saved to file in current 
	directory with the name 'res_CM.png'.

	Parameters
	----------
	true_classes : list of PyTorch IntTensors
		A list where each item is a [1 x len(sequence)] tensor with the true class
		label of the residues in a particular sequence
	predicted_classes : list of PyTorch FloatTensors
		A list where each item is a [1 x num_classes x len(sequence)] tensor
		with predictions of the class label for each residue in a particular
		sequence
	num_classes : int
		Number of distinct data classes
	output_dir : str, optional
		directory to which the plot will be saved (default is current directory)
	"""

	true_list = []
	pred_list = []

	for item in true_classes:
		single_frag = list(item[0].cpu().numpy().flatten())
		true_list = true_list + single_frag

	for item in predicted_classes:
		single_frag = item[0].permute(1, 0).cpu().numpy()

		for residue in single_frag:
			pred_list.append(np.argmax(residue))

	cm = np.zeros((num_classes, num_classes))
	for i in range(len(true_list)):
		cm[true_list[i], pred_list[i]] += 1

	df_cm = pd.DataFrame(cm, range(num_classes), range(num_classes))
	sn.set(font_scale=1.4) # for label size
	sn.heatmap(df_cm, cmap='Blues', annot=True, annot_kws={"size": 16}) # font size
	plt.ylabel('True labels')
	plt.xlabel('Predicted labels')
	plt.title('Test set confusion matrix')
	plt.tight_layout()
	plt.savefig(output_dir + 'res_CM.png')
	plt.show()

def output_predictions_to_file(sequence_data, excludeSeqID, encoding_scheme, encoder=None,
																			output_dir=''):
	"""Output sequences, their true values, and their predicted values to a file

	Used on the output of the test_unlabeled_data() function in the train_network module in
	order to detail the performance of the trained network on the test set. Produces the
	file "test_set_predictions.tsv" in output_dir. Each pair of lines in this tsvfile
	corresponds to a particular test set sequence, with the first containing the true data
	values, and the second line having the predicted data values.

	Parameters
	----------
	sequence_data : list of lists
		Details of the output predictions for each of the sequences in the test set. Each
		inner list represents a sample in the test set, with the format: [sequence_vector,
		true_value, predicted_value, sequence_ID]
	excludeSeqID : bool
		Boolean indicating whether or not each line in `tsvfile` has a sequence ID
		(default is False)		
	encoding_scheme : str
		Description of how an amino acid sequence should be encoded as a numeric 
		vector. Providing a string other than 'onehot', 'biophysics', or 'user' 
		will produce unintended consequences.
	encoder: UserEncoder object, optional
		If encoding_scheme is 'user', encoder should be a UserEncoder object
		that can convert amino acid sequences to numeric vectors. If
		encoding_scheme is not 'user', use None.
	"""

	seq_vectors = []
	true_vals = []
	pred_vals = []
	names = []
	count = 0
	for sequence in sequence_data:
		seq_vector, true_val, pred_val, name = sequence
		seq_vectors.append(seq_vector)
		true_vals.append(true_val)
		pred_vals.append(pred_val)

		if excludeSeqID:
			names.append('test' + str(count))
			count+=1
		else:
			names.append(name)

	# Decode the sequence vectors
	if encoding_scheme == 'onehot':
		sequences = encode_sequence.rev_one_hot(seq_vectors)
	elif encoding_scheme == 'biophysics':
		sequences = encode_sequence.rev_biophysics(seq_vectors)
	else:
		sequences = encoder.decode(seq_vectors)

	# Write to file
	filename = 'test_set_predictions.tsv'

	with open(output_dir + filename, 'w') as tsvfile:
		for i in range(len(names)):

			# Adjust formatting for residues or sequence data
			if isinstance(true_vals[i], np.ndarray):
				true_vals_format = ' '.join(true_vals[i].astype(str))
				pred_vals_format = ' '.join(pred_vals[i].astype(str))
			else:
				true_vals_format = true_vals[i]
				pred_vals_format = pred_vals[i]

			'''
			Format:
			NAME_TRUE SEQUENCE TRUE_VALUE(S)
			NAME_PRED SEQUENCE PRED_VALUE(S)
			'''
			output_str = "%s_TRUE %s %s\n" % (names[i], sequences[i], true_vals_format)
			output_str = output_str + "%s_PRED %s %s\n" % (names[i], sequences[i], pred_vals_format)
	
			tsvfile.write(output_str)
