import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras import activations, optimizers, layers, regularizers
import pandas
import os
import utils
import math
import re

class TextClassifier:
	"""
	Generalized classifier build on top of Keras Sequential() model.
	Created for classification of simple text, but should work on other data, as long as you respect it's limitations.
	"""
	#Default arguments for building and training model
	default_build_args = {
		'layers':[],
		'optimizer': optimizers.Adam(learning_rate=0.001),
		'model_name':'',
		'intermediate_dropout':False
		}

	default_train_args = {
		'training_data':None,
		'evaluation_data':None,
		'epochs':400
		}

	session = None
	graph = None
	model = None
	def __init__(self, build_args, train_args):
		"""
		Default values:

		build_args = {
		'layers':[],
		'optimizer':tf.keras.optimizers.Adam(),
		'model_name':'',
		'intermediate_dropout':False
		}

		train_args = {
		'training_data':None,
		'evaluation_data':None,
		'epochs':400
		}

		Data should be reshapable into square matrix with no padding or gaps. 
		Otherwise the convolution won't work.

		"""
		if build_args is None:
			self.default_build_args = default_build_args
		else:
			self.default_build_args = build_args
		if train_args is None:
			self.default_train_args = default_train_args
		else:
			self.default_train_args = train_args
	
	def train_model(self,
					training_dataset=None,
					validation_dataset=None,
					epochs=None,
					steps_per_epoch=None,
					model_name=None,
					write_gradients=True,
					write_images=False,
					class_weights=None
					):

		"""
		You can leave parameters empty if you have provided train_args dictionary to constructor,
		or override specific parameters for purposes of training.
		"""
		
		if training_dataset is None:
			if isinstance(self.default_train_args['training_data'], tf.data.Dataset):
				training_dataset = self.default_train_args['training_data']
			else:
				training_dataset = utils.data_to_tensors(dataset = self.default_train_args['training_data'],
														label_dictionary=self.default_train_args['label_dictionary'],
														feature_columns=self.default_train_args['data_column'],
														label_column=self.default_train_args['label_column'],
														batch_size=self.default_train_args.get('batch_size', 32),
														shuffle=self.default_train_args.get('shuffle', False))

		if validation_dataset is None:
			if isinstance(self.default_train_args['evaluation_data'], tf.data.Dataset):
				validation_dataset = self.default_train_args['evaluation_data']
			else:
				validation_dataset = utils.data_to_tensors(dataset=self.default_train_args['evaluation_data'],
														label_dictionary=self.default_train_args['label_dictionary'],
														feature_columns=self.default_train_args['data_column'],
														label_column=self.default_train_args['label_column'],
														batch_size=self.default_train_args.get('batch_size', 32),
														shuffle=self.default_train_args.get('shuffle', False))

		if epochs is None:
			epochs = self.default_train_args.get('epochs', 20)
		if steps_per_epoch is None:
			steps_per_epoch = self.default_train_args.get('steps_per_epoch', 50)
		if model_name is None:
			model_name = self.default_train_args.get('model_name',  
											self.default_build_args.get('model_name', str(os.getppid())))
		if class_weights is None:
			class_weights = self.default_train_args.get('class_weights', None)
			if class_weights is not None:
				class_weights = dict([(self.default_train_args['label_dictionary'][key], class_weights[key]) for key in class_weights.keys()])
		
		tensorboard_callback = keras.callbacks.TensorBoard(log_dir='.\models\\'+model_name, 
														histogram_freq=int(epochs/100), 
														write_grads=write_gradients,
														write_images=write_images)


		if validation_dataset is None:
			if isinstance(training_dataset, tf.data.Dataset):
				history = self.model.fit(training_dataset,
								epochs = epochs,
								steps_per_epoch = steps_per_epoch,
								shuffle = True,
								callbacks=[tensorboard_callback],
								class_weight=class_weights)
			else:
				history = self.model.fit(x=training_dataset[0],
								y=training_dataset[1],
								epochs = epochs,
								steps_per_epoch = steps_per_epoch,
								shuffle = True,
								callbacks=[tensorboard_callback],
								class_weight=class_weights)
		else:
			if isinstance(training_dataset, tf.data.Dataset):
				history = self.model.fit(training_dataset,
								validation_data = validation_dataset,
								epochs = epochs,
								steps_per_epoch = steps_per_epoch,
								validation_steps = 50,
								shuffle = True,
								callbacks=[tensorboard_callback],
								class_weight=class_weights)
			else:
				history = self.model.fit(x=training_dataset[0],
								y=training_dataset[1],
								validation_data = validation_dataset,
								epochs = epochs,
								steps_per_epoch = steps_per_epoch,
								validation_steps = 50,
								shuffle = True,
								callbacks=[tensorboard_callback],
								class_weight=class_weights)

		return history

	def build_model(self,
				dense_layers = None,
				convolution_layers = None,
				kernel_size = None,
				vocabulary_size = None,
				embedding_dimension = None,
				separable_convolution = None,
				locally_connected_convolution = None,
				max_pooling = None,
				pool_size = None,
				intermediate_dropout = None,
				intermediate_dropout_rate = None,
				activation = None,
				optimizer = None,
				dense_kernel_regulization_rate = None,
				dense_activity_regulization_rate = None,
				convolution_kernel_regulization_rate = None,
				name = None,
				input_size = None, 
				categories = None,
				exclusive_categories = None,
				input_l1_regularization = None,
				input_l2_regularization = None,
				clear_input = None,
				build_tensor_datasets = None):

		if dense_layers is None:
			dense_layers = self.default_build_args['layers']
		if convolution_layers is None:
			convolution_layers = self.default_build_args.get('convolution_layers', [])
		if kernel_size is None:
			kernel_size = self.default_build_args.get('kernel_size', 5)
		if vocabulary_size is None:
			vocabulary_size = self.default_build_args.get('vocabulary_size', 14000)
		if embedding_dimension is None:
			embedding_dimension = self.default_build_args.get('embedding_dim', 5)
		if separable_convolution is None:
			separable_convolution = self.default_build_args.get('separable_convolution', False)
		if locally_connected_convolution is None:
			locally_connected = self.default_build_args.get('locally_connected', False)
		if max_pooling is None:
			max_pooling = self.default_build_args.get('max_pooling', False)
		if pool_size is None:
			pool_size = self.default_build_args.get('pool_size', 10)
		if intermediate_dropout_rate is None:
			intermediate_dropout_rate = self.default_build_args.get('intermediate_dropout_rate', 0.0)
		if activation is None:
			activation = self.default_build_args.get('activation', 'relu')
		if optimizer is None:
			optimizer = self.default_build_args.get('optimizer', tf.keras.optimizers.Adam())
		if dense_kernel_regulization_rate is None:
			dense_kernel_regulization_rate = self.default_build_args.get('dense_l2_kernel_regulization_rate', None)
			if dense_kernel_regulization_rate is not None:
				dense_kernel_l2_regulizer = regularizers.L1L2(l2 = dense_kernel_regulization_rate)
			else:
				dense_kernel_l2_regulizer = None
		if dense_activity_regulization_rate is None:
			dense_activity_regulization_rate = self.default_build_args.get('dense_l2_activity_regulization_rate', None)
			if dense_activity_regulization_rate is not None:
				dense_activity_l2_regulizer = regularizers.L1L2(l2 = dense_activity_regulization_rate)
			else:
				dense_activity_l2_regulizer = None
		if convolution_kernel_regulization_rate is None:
			convolution_kernel_regulization_rate = self.default_build_args.get('convolution_kernel_l2_regulization_rate', None)
			if convolution_kernel_regulization_rate is not None:
				convolution_kernel_l2_regulizer = regularizers.L1L2(l2 = convolution_kernel_regulization_rate)
			else:
				convolution_kernel_l2_regulizer = None
		if name is None:
			name = self.default_build_args.get('model_name', '')
		if input_size is None:
			input_size = self.default_build_args.get('input_size', 100) 
		if categories is None:
			categories = self.default_build_args.get('categories', 10)
		if exclusive_categories is None:
			exclusive_categories = self.default_build_args.get('exclusive_categories', True)
		if input_l1_regularization is None:
			input_l1_regularization = self.default_build_args.get('input_l1regularization', 0.0)
		if input_l2_regularization is None:
			input_l2_regularization = self.default_build_args.get('input_l2regularization', 0.0)
		if clear_input is None:
			clear_input = self.default_build_args.get('clear_input', False)
		if build_tensor_datasets is None:
			build_tensor_datasets = self.default_build_args.get('build_tensor_datasets', False)

		
		#First build dataset tensors, if needed
		#This breaks separation between building and training. It's not a good design decision, but the benefits on functionality side are substantial.

		if build_tensor_datasets:
			self.default_train_args['training_data'] = utils.data_to_tensors(self.default_train_args['training_data'],
																	categories = categories,
																	noise = self.default_train_args.get('noise', 0.0), 
																	categorical_labels = self.default_train_args.get('categorical_labels', False),
																	data_type = self.default_train_args.get('data_type', 'text'),
																	label_column = self.default_train_args.get('label_column', None),
																	feature_columns = self.default_train_args.get('data_column', None),
																	shuffle = self.default_train_args.get('shuffle', None))

			self.default_train_args['evaluation_data'] = utils.data_to_tensors(self.default_train_args['evaluation_data'],
																	categories = categories,
																	noise = self.default_train_args.get('noise', 0.0),
																	categorical_labels = self.default_train_args.get('categorical_labels', False),
																	data_type = self.default_train_args.get('data_type', 'text'),
																	label_column = self.default_train_args.get('label_column', None),
																	feature_columns = self.default_train_args.get('data_column', None),
																	shuffle = self.default_train_args.get('shuffle', None))
			
		model = keras.Sequential()
		model.add(keras.layers.Embedding(input_dim=vocabulary_size, 
							   output_dim=embedding_dimension, 
							   input_length=input_size))

		##Filter layers
		for filter in convolution_layers:
			model.add(tf.keras.layers.Convolution1D(filters=filter, kernel_size=kernel_size))
			if pool_size > 0:
				if max_pooling:
					model.add(tf.keras.layers.MaxPooling1D(pool_size=pool_size, 
											strides=None, 
											padding='valid', 
											data_format='channels_last'))
				else:
					model.add(tf.keras.layers.AveragePooling1D(pool_size=pool_size, 
												strides=None, 
												padding='valid', 
												data_format='channels_last'))
		
		model.add(layers.Flatten())
			
		for layer in dense_layers:
			model.add(tf.keras.layers.Dense(layer, activation=activation, 
								   kernel_regularizer=dense_kernel_l2_regulizer,
								   activity_regularizer=dense_activity_l2_regulizer))
			#Dropout layer
			if intermediate_dropout_rate > 0.0:
				model.add(tf.keras.layers.Dropout(intermediate_dropout_rate))
			#Normalization 
			#Maybe move it before dropout, also added trainable flag
			if self.default_build_args.get('normalization', False):
				model.add(tf.keras.layers.BatchNormalization(
					momentum=self.default_build_args.get('normalization_momentum', 0.999),
					trainable=True))

		if exclusive_categories:
			# Compute logits (1 per class).
			model.add(tf.keras.layers.Dense(categories, activation='softmax'))
	
			#Initiating optimizer
			model.compile(optimizer=optimizer, 
					loss='categorical_crossentropy', 
					metrics=['accuracy', 'categorical_accuracy', keras.metrics.Recall(), keras.metrics.Precision()])
		else:
			# Compute logits (1 per class).
			model.add(tf.keras.layers.Dense(categories, activation='sigmoid'))
	
			#Initiating optimizer
			model.compile(optimizer=optimizer, 
					loss='binary_crossentropy', 
					metrics=['accuracy','categorical_accuracy', keras.metrics.Recall(), keras.metrics.Precision()])

		model.summary()
		self.model = model

		return model

	def test_model(self, testing_dataset, vocabulary_dict, label_dictionary, batch_size, steps, print_desc=False):
		"""
		testing_dataset - (data, labels) tuple. As provided by construct_simple_dataset function.

		"""

		results = self.model.predict(testing_dataset[0], batch_size=batch_size, verbose=0)

		evaluation_labels = testing_dataset[1]

		testing_dataset = testing_dataset[0]

		label_dictionary = dict((label_dictionary[element], element) for element in label_dictionary)

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

		def categories_to_string(array):
			if array.size == 0:
				return ""
			for i, element in np.ndenumerate(array):			
				element = label_dictionary[i]			
			return array

		results = np.rint(results)
		results = results.astype(int)
		correct = 0

		#Printing results
		for i, element in enumerate(results):
			if (element==evaluation_labels[i]).all():
				correct = correct+1
			elif print_desc:
				print(utils.tokens_to_text(description_text=testing_dataset[i], 
								vocabulary_dict=vocabulary_dict) + '\n\n' +
								"Predicted: " + str(utils.decode_label(element, label_dictionary)) +
								" Expected: " + str(utils.decode_label(evaluation_labels[i], label_dictionary)) + '\n\n\n')
			else:
				print("Predicted: ",label_dictionary[element]," Expected: " ,label_dictionary[evaluation_labels[i]])
		print('Final Accuracy', correct/results.shape[0])
		return results	