import pandas as pd
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import random
import os
import tensorflow.keras as keras
from tensorflow.keras import backend as K

import classifiers
import utils

def plot_embedding(model, 
				   informative_terms):
	embedding_matrix = model.get_variable_value('input_layer/risk_desc_embedding/embedding_weights')

	fig = plt.figure(dpi=300)
	for term_index in range(300):
	  # Create a one-hot encoding for our term. It has 0s everywhere, except for
	  # a single 1 in the coordinate that corresponds to that term.
	  term_vector = np.zeros(len(informative_terms))
	  term_vector[term_index] = 1
	  # We'll now project that one-hot vector into the embedding space.
	  embedding_xy = np.matmul(term_vector, embedding_matrix)
	  plt.text(embedding_xy[0],
			   embedding_xy[1],
			   informative_terms[term_index])

	# Do a little setup to make sure the plot displays nicely.
	plt.rcParams["figure.figsize"] = (150, 150)
	plt.xlim(1.2 * embedding_matrix.min(), 1.2 * embedding_matrix.max())
	plt.ylim(1.2 * embedding_matrix.min(), 1.2 * embedding_matrix.max())
	fig.savefig('./figs/embeddingplot'+str(os.getpid())+str(random.randint(0,999))+'.png')

def plot_loss(loss_data_lines, 
			  accuracy_data_lines, 
			  final_performance_points):
	fig = plt.figure(dpi=600)
	plt.title("Average loss during training")

	plt.subplot(211)
	plt.ylabel("Avg. loss")
	for n, line in enumerate(loss_data_lines):
		plt.plot(line, label="n={0}".format(n))
	plt.legend(title="Legend")
	plt.subplot(212)

	plt.ylabel("Accuracy")
	plt.xlabel("epoch")
	for n, line in enumerate(accuracy_data_lines):
		plt.plot(line, label="n={0}".format(n))
	plt.legend(title="Legend")
	print(final_performance_points)
	fig.savefig('./figs/'+str(os.getpid())+str(random.randint(0,999))+'.png')

def create_label_file(vocabulary):
	with open('./vocabulary.tsv','w', encoding="utf-8") as tsvFile:
		for word in vocabulary:
			tsvFile.write(word+"\n")
	tsvFile.close()

def safe_model(model, path):

	feature_spec = {'risk_desc' : tf.FixedLenFeature(shape=[1], dtype=tf.string, default_value=None) }
	
	serving_input_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(
								feature_spec)
	model.export_savedmodel(path, serving_input_fn)

def main():
	#Message level, to keep console tidy
	os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
	training_dataset, \
		evaluation_dataset, \
		testing_dataset, \
		word_dictionary, \
		label_dictionary, \
		word_export_vocabulary, \
		max_desc_len = utils.import_data_text('./data/rapexdata.csv',
												feature_columns = ['risk_desc'],
												strict_sanitization=True, 
												matrix_form=True, 
												tokenize_words=True, 
												keep_composite_labels=True,
												merge_input_cols=True,
												to_tensors=True)

	all_loss_data = []
	all_accuracy_data = []
	#Hyperparams
	steps = 50
	batch_size = 64

	build_args = {
		'layers':[200,100,50],
		'convolution_layers':[128,512],
		'embedding_dim':5,
		'optimizer':keras.optimizers.Adam(),
		'model_name':'alpha',
		'vocabulary_size':len(word_dictionary)+1,
		'kernel_size':10,
		'pool_size':5,
		'max_pooling':True,
		'exclusive_categories': False,
		'input_size':max_desc_len,
		'categories':len(label_dictionary)
		}

	build_argsb = {
		'layers':[2000, 500, 400, 300, 200, 100, 50],
		'convolution_layers':[128,512],
		'embedding_dim':5,
		'optimizer':keras.optimizers.Adam(),
		'model_name':'beta',
		'vocabulary_size':len(word_dictionary)+1,
		'kernel_size':10,
		'pool_size':5,
		'max_pooling':True,
		'exclusive_categories': False,
		'input_size':max_desc_len,
		'categories':len(label_dictionary)
		}
	build_argsc = {
		'layers':[200,100,50],
		'convolution_layers':[128,512],
		'embedding_dim':20,
		'optimizer':keras.optimizers.Adam(),
		'model_name':'gamma',
		'vocabulary_size':len(word_dictionary)+1,
		'kernel_size':10,
		'pool_size':5,
		'max_pooling':True,
		'exclusive_categories': False,
		'input_size':max_desc_len,
		'categories':len(label_dictionary)
		}
	build_argsd = {
		'layers':[2000, 500, 400, 300, 200, 100, 50],
		'convolution_layers':[128,512],
		'embedding_dim':20,
		'optimizer':keras.optimizers.Adam(),
		'model_name':'delta',
		'vocabulary_size':len(word_dictionary)+1,
		'kernel_size':10,
		'pool_size':5,
		'max_pooling':True,
		'exclusive_categories': False,
		'input_size':max_desc_len,
		'categories':len(label_dictionary)
		}
	train_args = {
		'training_data':training_dataset,
		'evaluation_data':evaluation_dataset,
		'epochs':100,
		'steps_per_epoch':steps,
		'batch_size':batch_size,
		'label_dictionary':label_dictionary,
		'data_column':'risk_desc',
		'label_column':'risk_type',
		'shuffle':True
		}

	model = classifiers.TextClassifier(build_args, train_args)

	model.build_model()
	model.train_model()
	model.test_model(testing_dataset, 
				  word_export_vocabulary, 
				  label_dictionary, 
				  batch_size, 
				  steps=steps, 
				  print_desc=True)
	tf.keras.backend.clear_session()

	modelb = classifiers.TextClassifier(build_argsb, train_args)

	modelb.build_model()
	modelb.train_model()
	modelb.test_model(testing_dataset, 
				  word_export_vocabulary, 
				  label_dictionary, 
				  batch_size, 
				  steps=steps, 
				  print_desc=True)
	tf.keras.backend.clear_session()

	modelc = classifiers.TextClassifier(build_argsc, train_args)

	modelc.build_model()
	modelc.train_model()
	modelc.test_model(testing_dataset, 
				  word_export_vocabulary, 
				  label_dictionary, 
				  batch_size, 
				  steps=steps, 
				  print_desc=True)
	tf.keras.backend.clear_session()

	modeld = classifiers.TextClassifier(build_argsd, train_args)

	modeld.build_model()
	modeld.train_model()
	modeld.test_model(testing_dataset, 
				  word_export_vocabulary, 
				  label_dictionary, 
				  batch_size, 
				  steps=steps, 
				  print_desc=True)


if __name__ == '__main__':
	main()