import numpy as np
from scipy import ndimage, misc
from glob import glob
import re
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras import activations, utils, layers
import pandas as pd
import random
import math
from multiprocessing import Pool
from matplotlib import pyplot
import os

###################################
###HELPER FUNCTIONS################
# Used in other utility functions.#
###################################

def get_max_string_feature_length(series):
	max_len = 0
	for element in series:
		#If we already have list as input we don't need to split it
		if type(element) is list:
			max_len = max(max_len, len(element))
		else:
			max_len = max(max_len, len(element.split(' ')))
	return max_len

#Self explanatory
def get_average_string_feature_length(series):
	avg_len = 0
	if series[0] is np.array:
		avg_len = np.average(series.apply(lambda x: x.size))
	else:
		avg_len = np.average(series.apply(lambda x: len(x.split(' '))))
	return avg_len


def byte_to_float(value):
	return value/255

def add_noise(input_values, noise_intensity, max_value=1.0, min_value=0.0):

	noise = np.random.normal(size = input_values.size, scale = noise_intensity)
	input_values = np.add(input_values, noise)
	input_values = np.absolute(input_values)
	input_values = np.clip(input_values, a_max=max_value, a_min=min_value)

	return input_values

def encode_labels(label, label_dictionary):
		new_label = [0] * len(label_dictionary)
		if type(label) is list:
			for element in label:
				new_label[element] = 1
		else:
			new_label[label] = 1
		return new_label

def create_label_dictionary(label_series):
	#Find all label values
	label_values=[]
	if label_series.dtype == 'O':
		for element in pd.unique(label_series):
			label_values.extend([re.sub('^ ','',value) for value in element.split(',')])
	else:
		for element in pd.unique(label_series):
			label_values.append(element)

	label_values = enumerate(list(set(label_values)))

	#put them in dict
	label_dictionary = dict((element,i) for i, element in label_values)

	return label_dictionary

###################################
#End of helper functions.         #
###################################

def import_numerical_data_csv(training_data_path, 
					testing_data_path=None,
					feature_columns = ['risk_desc','defect_desc'],
					label_column = 'risk_type',
					categorical_labels = True,
					dataset_split_denominator=5,
					to_tensors=True, 
					training_noise_intensity=0.0):
	"""
	Reads numerical data from csv files. Adds noise and converts to tensors if needed.

	training_data_path - unix like path to training data, will be split into training and evaluation datasets
	testing_data_path - unix like path to testing data
	dataset_split_denominator - denominator for splitting dataset into training and evaluation, 5 by default
	to_tensors - convert data to tensor datasets, default True
	training_noise_intensity - the scale of random noise added to the training sample, requires conversion to tensors.

	Returns training_dataset, evaluation_dataset and testing_dataset

	"""
	if testing_data_path is None:
		print("No testing dataset path provided. Spliting the training dataset with default parameters.")
		primary_dataset = pd.read_csv(training_data_path)
		training_dataset, evaluation_dataset, testing_dataset = split_dataset(primary_dataset, 
																		three_fold = True,  
																		dataset_split_denominator = dataset_split_denominator)
	else:
		primary_dataset = pd.read_csv(training_data_path)
		testing_dataset = pd.read_csv(testing_data_path)
		
		training_dataset = primary_dataset[int(primary_dataset.shape[0]/dataset_split_denominator):]
		evaluation_dataset = primary_dataset[:int(primary_dataset.shape[0]/dataset_split_denominator)]

	#Remove all columns except for those listed in feature_columns and label column
	training_dataset = training_dataset[feature_columns+[label_column]]
	evaluation_dataset = evaluation_dataset[feature_columns+[label_column]]
	testing_dataset = testing_dataset[feature_columns+[label_column]]

	if to_tensors:
		training_dataset = data_to_tensors(training_dataset, 
									 noise = training_noise_intensity,
									 label_column = label_column,
									 categorical_labels = categorical_labels)
		evaluation_dataset = data_to_tensors(evaluation_dataset, 
									 categorical_labels = categorical_labels)

	return training_dataset, evaluation_dataset, testing_dataset

def import_data_image(training_data_path, 
					  testing_data_path, 
					  training_labels_path,
					  feature_columns = ['risk_desc','defect_desc'],
					  label_column = 'label',
					  categories = 10, 
					  testing_labels_path=None,
					  file_name_pattern = '*.jpg',
					  to_tensors=True, 
					  training_noise_intensity=0.0):
	"""
	Reads data from csv files. Adds noise and converts to tensors if needed.

	training_data_path - unix like path to training data, will be split into training and evaluation datasets
	testing_data_path - unix like path to testing data
	testing_labels_path - unix like path to testning data
	categories - number of categories for classification
	to_tensors - convert data to tensor datasets, default True
	training_noise_intensity - the scale of random noise added to the training sample, requires conversion to tensors.

	Returns training_dataset, evaluation_dataset and testing_dataset

	"""
	training_files_list = glob(training_data_path+'*')
	testing_files_list = glob(testing_data_path+'*')
	

	primary_dataset = np.array([np.array((path.split(sep='\\')[1], misc.imread(path))) for path in training_files_list])
	testing_dataset = np.array([np.array((path.split(sep='\\')[1], misc.imread(path))) for path in testing_files_list])

	columns = ['id', 'data']
	primary_dataset = pd.DataFrame(primary_dataset, columns = columns)
	testing_dataset = pd.DataFrame(testing_dataset, columns = columns)

	primary_labels = pd.read_csv(training_labels_path)

	if testing_labels_path is not None:
		testing_labels = pd.read_csv(testing_labels_path)

	primary_dataset = primary_dataset.join(primary_labels.set_index('id'), on='id')

	primary_dataset = primary_dataset.drop(['id'], axis=1)

	label_dictionary = create_label_dictionary(primary_dataset[label_column])
	
	training_dataset = primary_dataset[int(primary_dataset.shape[0]/5):]
	evaluation_dataset = primary_dataset[:int(primary_dataset.shape[0]/5)]

	if to_tensors:
		training_dataset = data_to_tensors(training_dataset, 
									 label_column = label_column,
									 label_dictionary=label_dictionary,
									 noise = training_noise_intensity)

		evaluation_dataset = data_to_tensors(evaluation_dataset,
									 label_column = label_column, 
									 label_dictionary=label_dictionary,
									 noise = training_noise_intensity)

	return training_dataset, evaluation_dataset, testing_dataset

def import_data_text(training_data_path,
					testing_data_path = None,
					feature_columns = ['risk_desc','defect_desc'],
					label_column = 'risk_type',
					categorical_labels = True,
					main_feature_column = None,
					max_feature_len=0,
					keep_composite_labels=False, 
					strict_sanitization=False, 
					tokenize_words=False, 
					matrix_form=False,
					training_noise_intensity=0.0,
					merge_input_cols=False,
					drop_columns=[],
					categorical_weights = None,
					to_tensors=True):
	"""
	Reads data from csv files. Adds noise and converts to tensors if needed.

	training_data_path - unix like path to training data, will be split into training and evaluation datasets
	testing_data_path - unix like path to testing data
	to_tensors - convert data to tensor datasets, default True
	training_noise_intensity - the scale of random noise added to the training sample, requires conversion to tensors.

	Returns training_dataset, evaluation_dataset and testing_dataset

	"""
	
	#Counts number of label elements
	def label_lenght(label):
		return len(label.split(','))

	#Returns the first tag 
	def first_label(label):
		return label.split(',')[0]

	def composite_label_translation(label):
		result = [label_dictionary[re.sub('^ ','',element)] 
		  for element in label.split(',')]
		return result
	
	def tokenize(element):
		newelement = ''
		for word in element.split(' '):
			newelement += ' ' + str(word_dictionary[word])
		newelement = newelement[1:]
		return newelement

	def to_matrix(element):
		array_form = element.split(' ')
		if tokenize_words:
			array_form = [int(i) for i in array_form]
		return np.array(array_form)

	def encode_text(cell):
		if len(cell.split(' ')) < max_feature_len:
			cell = cell.split(' ')
			cell.extend([padding]*(max_feature_len-len(cell)))
			cell = ' '.join(cell)
		if len(cell.split(' ')) > max_feature_len:
			cell = cell.rsplit(' ',  len(cell.split(' '))-max_feature_len)[0]
		if tokenize_words:
			testvar=len(cell.split(' '))
			cell = tokenize(cell)
		if matrix_form:
			testvar=len(cell.split(' '))
			cell = to_matrix(cell)
		return cell

	if testing_data_path is None:
		print("No testing dataset path provided. Spliting the training dataset with default parameters.")
		primary_dataset = pd.read_csv(training_data_path)
	else:
		primary_dataset = pd.read_csv(training_data_path)
		testing_dataset = pd.read_csv(testing_data_path)	

	if len(drop_columns) > 0:
		primary_dataset.drop(drop_columns, axis=1, inplace=True)

	padding = '<PAD>'

	for feature_column in feature_columns:
		primary_dataset[feature_column] = clean_text(primary_dataset[feature_column],
								 strict_sanitization=strict_sanitization)	
	
	word_list = []
	#adding padding symbol
	word_list.append(padding)	
	
	# and we get rid of all composite labels
	if not keep_composite_labels:
		print("Removing composite labels.")
		primary_dataset = primary_dataset[primary_dataset[label_column].map(label_lenght) == 1]

	#put them in dict
	label_dictionary = create_label_dictionary(primary_dataset[label_column])

	#We replace the risks in data with their dictionary values
	if not keep_composite_labels:
		primary_dataset[label_column] = primary_dataset[label_column].map(lambda element:label_dictionary[element.split(',')[0]])
	else:
		primary_dataset[label_column] = primary_dataset[label_column].map(composite_label_translation)

	if main_feature_column is None:
		main_feature_column = feature_columns[0]

	#Merging columns
	if merge_input_cols:
		for feature_column in feature_columns[1:]:
			primary_dataset[main_feature_column] = primary_dataset[main_feature_column] + primary_dataset[feature_column]
		primary_dataset = primary_dataset.drop(columns=feature_columns[1:])

	#Removing empty descriptions
	primary_dataset = primary_dataset.dropna(subset=[main_feature_column])

	#Shuffling dataset
	primary_dataset = primary_dataset.sample(frac=1.0)

	for element in primary_dataset[main_feature_column]:
			word_list.extend(element.split(' '))
	
	#Get rid of duplicate words
	word_list = list(set(word_list))
	#We create dictionary of all words used in descriptions
	word_dictionary = dict((element,i) for i, element in enumerate(word_list))

	word_export_vocabulary = dict((i,element) for i, element in enumerate(word_list))

	if max_feature_len is 0:
		max_feature_len = get_max_string_feature_length(primary_dataset[main_feature_column])
	average_feature_len = get_average_string_feature_length(primary_dataset[main_feature_column])
	primary_dataset[main_feature_column] = primary_dataset[main_feature_column].map(encode_text)

	if testing_data_path is None:
		training_dataset, evaluation_dataset, testing_dataset = split_dataset(primary_dataset, 
																		three_fold = True)

	print("\n\n\nMax text lenght:", 
	   max_feature_len,
	   "\nAverage text lenght:",
		average_feature_len,
	   "\nDataset size:", 
	   primary_dataset[main_feature_column].size, 
	   "\nUnique words:", len(word_list), 
	   "\nNumber of labels:", len(label_dictionary))

	if to_tensors:
		training_dataset = data_to_tensors(training_dataset, 
									 label_dictionary = label_dictionary, 
									 noise = training_noise_intensity,
									 label_column = label_column, 
									 feature_columns = main_feature_column,
									 categorical_labels = categorical_labels,
									 categorical_weights = categorical_weights)
		evaluation_dataset = data_to_tensors(evaluation_dataset,
									 label_dictionary = label_dictionary,
									 label_column = label_column, 
									 feature_columns = main_feature_column,
									 categorical_labels = categorical_labels,
									 categorical_weights = categorical_weights)

	testing_dataset = construct_simple_dataset(testing_dataset, 
											label_dictionary, 
											label_column=label_column, 
											feature_columns=main_feature_column, 
											tokenized_data=tokenize_words)
	
	return training_dataset, evaluation_dataset, testing_dataset, word_dictionary, label_dictionary, word_export_vocabulary, max_feature_len

def split_dataset(primary_dataset, 
				  three_fold=False, 
				  dataset_split_denominator = 5):

	if three_fold:

		testing_dataset = primary_dataset[int(primary_dataset.shape[0]/dataset_split_denominator):2*int(primary_dataset.shape[0]/dataset_split_denominator)]

		training_dataset = primary_dataset[2*int(primary_dataset.shape[0]/dataset_split_denominator):]

		evaluation_dataset = primary_dataset[:int(primary_dataset.shape[0]/dataset_split_denominator)]

		return training_dataset, evaluation_dataset, testing_dataset
	else:
		training_dataset = primary_dataset[int(primary_dataset.shape[0]/dataset_split_denominator):]

		evaluation_dataset = primary_dataset[:int(primary_dataset.shape[0]/dataset_split_denominator)]

		return training_dataset, evaluation_dataset

def clean_text(series, 
			   strict_sanitization=False, 
			   append_space=True):

	def sanitize_text(cell, append_space=False):	
		#We check if the description is string
		if type(cell) is not str:
			# If not we convert to string
			cell = str(cell)
		#Next we sanitize the text
		#First we remove html tags etc.
		cell = re.sub(r'\<.*\>','',cell)
		if strict_sanitization:
			#Now we get rid of special chars
			cell = re.sub(r'[\.\,\'\"\?\!\)\(\>\<\/\_\{\}]',' ',cell)
			#Get rid of mathematical operators, replace with space
			cell = re.sub(r'[\+\-\*\%\:\°]',' ',cell)
			#and once again few specials
			cell = re.sub(r'[\<\>\;]','',cell)
		# get everythin in lowercase
		cell = cell.lower()
		if append_space:
			cell = cell + ' '
		return cell
	
	series = series.map(sanitize_text)

	return series

def exploratory_analysis(dataset):
	"""
	Describe selected dataset.
	"""
	print(dataset.describe())

def plot_matrix(cm):
	"""
	Plots confusion matrix using pyplot. 
	Shamelessly lifted from https://scikit-learn.org/stable/auto_examples/model_selection/plot_confusion_matrix.html
	with minimal changes.

	"""
	fig, ax = pyplot.subplots()
	im = ax.imshow(cm, interpolation='nearest', cmap=pyplot.cm.Blues)
	ax.figure.colorbar(im, ax=ax)
	classes = [str(i) for i in range(0,cm.shape[0])]
	# We want to show all ticks...
	ax.set(xticks=np.arange(cm.shape[1]),
			yticks=np.arange(cm.shape[0]),
			# ... and label them with the respective list entries
			xticklabels=classes, yticklabels=classes,
			title="Confusion_matrix",
			ylabel='True label',
			xlabel='Predicted label')

	# Rotate the tick labels and set their alignment.
	pyplot.setp(ax.get_xticklabels(), rotation=45, ha="right",
				rotation_mode="anchor")

	# Loop over data dimensions and create text annotations.
	fmt = '.2f'
	thresh = cm.max() / 2.
	for i in range(cm.shape[0]):
		for j in range(cm.shape[1]):
			ax.text(j, i, format(cm[i, j], fmt),
					ha="center", va="center",
					color="white" if cm[i, j] > thresh else "black")
	fig.tight_layout()
	pyplot.show()

def plot_history(history):

	categorical_accuracy = history[1].history['val_categorical_accuracy']

	pyplot.plot(history[1].epoch, categorical_accuracy)

	pyplot.show()

def confusion_matrix(predictions, dims=(10,10)):

	def remove_nans(value):
		if math.isnan(value):
			return 0
		return value

	vectorized = np.vectorize(remove_nans)

	matrix = np.zeros(dims)

	dims = predictions[0][0].shape[1]

	for prediction in predictions:
		matrix[prediction[1]] = np.add(prediction[0][0], matrix[prediction[1]])

	row_sums = matrix.sum(axis=1)
	matrix = matrix / row_sums[:, np.newaxis]
	matrix = vectorized(matrix) # To remove all nans

	return matrix

def generate_individual(max_width, max_height):

	individual = []

	for layer in range(0, random.randint(1, max_height)):
		individual.append(16*random.randint(1,(max_width/16)))

	return individual

def generate_population(size, 
						max_width = 512,
						max_height = 5):

	population = []

	for individual in range(0,size):

		population.append(generate_individual(max_width, max_height))

	return population

def mutate_genome(individual, 
				  mutation_prob, 
				  max_height=5, 
				  max_width=512):

	for gene in individual:	
		if random.random() <= mutation_prob:
			gene = abs(
				gene + 
				(16*random.randint(-1,1))
				)%512

	if (random.random() < mutation_prob*mutation_prob 
	 and len(individual) > 1 
	 and len(individual) < max_height):
		if random.random()> 0.5:
			individual.append(16)

		else:
			individual.pop()

	return individual

def data_to_tensors(dataset, 
					label_dictionary = None,
					noise = 0.0, 
					categorical_labels = True, 
					categorical_weights = None,
					batch_size = 32,
					shuffle = True,
					label_column = 'label', 
					feature_columns = 'data',
					session = tf.compat.v1.Session()):
	"""
	Converts raw data to tensors.
	"""

	#Data are already shuffled no need to do it again.		
	even_data = dataset

	#Separate Labels from data
	labels = dataset[label_column]

	if categorical_labels:
		labels = labels.apply(encode_labels, args = (label_dictionary,))
		labels = np.array(labels.to_list(), np.int32)
	else:
		labels = np.array(labels.to_list(), np.float32)

	#We get rid of label column
	dataset.drop(label_column, axis=1, inplace=True)

	dataset = np.array(dataset[feature_columns].to_list(),np.int32)
	
	# I tried keeping the string format, now I'm trying to change the type of the ndarray.
	if categorical_weights is not None:
		weights = generate_sample_weights(labels, categorical_weights, label_dictionary)
		weights = np.array(weights, dtype=np.float32)
		dataset = tf.data.Dataset.from_tensor_slices((dataset, labels, weights))
	else:
		dataset = tf.data.Dataset.from_tensor_slices((dataset, labels))

	dataset = dataset.batch(batch_size)

	if shuffle:
		dataset = dataset.shuffle(int(even_data.shape[0]))

	dataset = dataset.repeat()

	return dataset

def construct_simple_dataset(dataset,
							label_dictionary,
							label_column = 'label', 
							feature_columns = 'data',
							categorical_labels = True,
							tokenized_data = True):

	#Data are already shuffled no need to do it again.		
	even_data = dataset

	#one hot encoding of lables, returns n-tuple with k elements valued 1, where n is total number of labels and k is number of fitting labels
	

	#Separate Labels from data
	labels = dataset[label_column].to_frame()

	if categorical_labels:
		labels = labels[label_column].apply(encode_labels, args = (label_dictionary,)).to_frame()	
		labels = np.array(labels[label_column].to_list(), np.int32)
	else:
		labels = np.array(labels[label_column].to_list(), np.float32)

	#We get rid of label column
	dataset.drop(label_column, axis=1, inplace=True)

	if tokenized_data:
		dataset = np.array(dataset[feature_columns].to_list(),np.int32)

	return (dataset, labels)

def generate_sample_weights(labels, weights, label_dictionary):
	new_weights = []
	#Recalculate weights
	weights = dict([(label_dictionary[key], weights[key]) for key in weights.keys()])

	for label in labels:
		label_weight = [1.0]*len(label_dictionary)

		for key in weights.keys():
			label_weight[key] = weights[key]

		new_weights.append(label_weight)

	return new_weights

def one_hot_to_labels(encoded_labels):
	"""
	Converts encoded labels into decimal equivalents. 
	Returns list of decoded labels.
	"""
	decoded_labels = []
	for label in encoded_labels:
		decoded_labels.append(np.argmax(label, axis=0))

	return decoded_labels

def write_predictions(labels, 
					  path = './predictions/new.csv'):
	"""
	Writes predicted labels into csv, by default in './predictions/new.csv'.
	With following format:

	ImageId,Label
	1,3
	2,7
	3,8 

	Note that the labels start with 1. 
	"""
	#Dictionary of sample IDs and predicted labels, ID's start from 1
	predictions = {
		'ImageId': [i for i in range(1, len(labels)+1)],
	   'Label':labels}

	df = pd.DataFrame(data=predictions)

	df.to_csv(path_or_buf = path, index = False)

	return df

def decode_label(encoded_label, 
				 label_dictionary):
	new_array = []
	for i, element in np.ndenumerate(encoded_label):
		if element == 1:
			new_array.append(label_dictionary[i[0]])
	
	return new_array


def tokens_to_text(description_text, 
				   vocabulary_dict):
	#Decoding description
	#element = decode_label(predicted_label, label_dictionary)
	#true_labels = decode_label(evaluation_label, label_dictionary)
	description = ' '.join([vocabulary_dict[coded_word] for coded_word in description_text])
	text = re.sub(r'\<PAD\>', '', description)
	return text