#!/usr/bin/env python

import sys, os
import argparse
import json
import csv

from sklearn.cross_validation import train_test_split
from keras import layers
from keras.callbacks import History, EarlyStopping
from keras import optimizers
from keras.objectives import binary_crossentropy

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

from data.utils import load_sparse_csr
from analysis import visualization, evaluate
from models import FullConnected, Convolution, TimeBased

BATCH_SIZE = 32

def log_result(train_accuracy, test_accuracy, auc, tpr, fpr, args):
	data_point = [args.architecture, args.params, args.embedding, train_accuracy, test_accuracy, auc, tpr, fpr]
	with open('./results.csv', 'a') as f:
		writer = csv.writer(f)
		writer.writerow(data_point)

def parse_optimizer(args):
	"""
		Parses the optimizer to use
	"""
	if args.optimizer == 'adam':
		return optimizers.Adam(args.learningrate)
	elif args.optimizer == 'sgd':
		return optimizers.SGD(args.learningrate)
	else:
		raise ValueError('Unknown optimizer')
	

def load_data(args):
	"""
		Loads the dataset according to the arguments
	"""
	embedding = args.embedding
	FOLDER = None

	if 'tfidf' in embedding or 'counts' in embedding:
		ngram1 = args.ngram_min
		ngram2 = args.ngram_max
		vocabsize = args.vocabsize
		if vocabsize not in [5000, 10000]:
			raise ValueError('Unknown vocab size')
		FOLDER = './data/processed/tfidf_'+str(vocabsize)+'_'+str(ngram1)+'-'+str(ngram2)+'grams'

		print('loading data from:' + FOLDER)

		X = load_sparse_csr(FOLDER + '/X_train.csr.npz').toarray()
		Y = np.load(FOLDER + '/Y_train.npy')


	elif 'glove' in embedding or 'learned' in embedding:
		embeddingsize = args.embeddingsize
		vocabsize = args.vocabsize
		if vocabsize not in [5000, 9000, None, 500]:
			raise ValueError('Unknown vocab size')
		if embeddingsize not in [100, 50]:
			raise ValueError('Unknown Embedding size')
		FOLDER = './data/processed/glove'+str(embeddingsize)+'_top'+str(vocabsize)

		print('Loading data from: ' + FOLDER)
		X = np.load(FOLDER + '/X_train.npy')
		Y = np.load(FOLDER + '/Y_train.npy')
	elif 'discwords' in embedding:
		embeddingsize = args.embeddingsize
		if embeddingsize not in [100, 50]:
			raise ValueError('Unknown Embedding size')
		FOLDER = './data/processed/discwords/glove'+str(embeddingsize)

		print('Loading data from: ' + FOLDER)
		X = np.load(FOLDER + '/X_train.npy')
		Y = np.load(FOLDER + '/Y_train.npy')
	else:
		raise IOError('Unknown embedding ' + embedding + ' supplied.')

	return train_test_split(X, Y, test_size=0.1), FOLDER

def main(args):
	"""
		Loading data
	"""
	params = json.load(open(args.params, 'r')) # information about architecture
	(X, X_val, Y, Y_val), folder = load_data(args) # dataset to use
	
	if args.use_test:
		X = np.vstack([X, X_val])
		Y = np.hstack([Y, Y_val])
		try:
			X_val = np.load(folder + '/X_test.npy')
		except Exception as e:
			X_val = load_sparse_csr( folder + '/X_test.csr.npz').toarray()
			
		Y_val = np.load(folder + '/Y_test.npy')

	if args.subset:
		X = X[:args.subset, :]
		Y = Y[:args.subset]

	FEATURE_SIZE = X.shape[1]

	"""
		Model building
	"""
	model = None
	if args.embedding not in ['glove', 'learned', 'discwords']:
		if args.architecture == 'FNN':
			model = FullConnected.FC(FEATURE_SIZE, **params)
		elif args.architecture == 'CNN':
			model = Convolution.CNN(FEATURE_SIZE, **params)
		else:
			raise ValueError('Unsupported architecture ' + args.architecture)
	else:
		"""
			Glove Embeddings!
		"""
		embedding = None
		
		if args.embedding == 'learned':
			LEN_WORD_MAP = np.max(np.vstack([X, X_val ])) #maximimum embedding size
			embedding = layers.Embedding(LEN_WORD_MAP + 1, args.embeddingsize, input_length=FEATURE_SIZE)

		elif args.embedding == 'glove' or args.embedding == 'discwords':
			embedding_matrix = np.load(folder + '/embedding_matrix.npy')
			LEN_WORD_MAP = embedding_matrix.shape[0]
			EMBED_SIZE = embedding_matrix.shape[1]

			embedding = layers.Embedding( LEN_WORD_MAP, EMBED_SIZE, \
				input_length=FEATURE_SIZE, weights=[embedding_matrix], \
				trainable=args.embedding_trainable)

		else:
			raise ValueError('Unsupported embedding')

		if args.architecture == 'FNN':
			model = FullConnected.FC_embedding(FEATURE_SIZE, embedding=embedding, **params)
		elif args.architecture == 'CNN':
			model = Convolution.CNN_embedding(FEATURE_SIZE, embedding=embedding, **params)
		elif args.architecture == 'LSTM':
			model = TimeBased.LSTM_embedding(FEATURE_SIZE, embedding=embedding, **params)
		elif args.architecture == 'LSTMCNN':
			model = TimeBased.LSTMCNN_embedding(FEATURE_SIZE, embedding=embedding, **params)
		elif args.architecture == 'biLSTMCNN':
			## compatability for saving results only...
			model = TimeBased.LSTMCNN_embedding(FEATURE_SIZE, embedding=embedding, **params)
		elif args.architecture == 'biLSTM':
			model = TimeBased.biLSTM_embedding(FEATURE_SIZE, embedding=embedding, **params)	
		else:
			raise ValueError('Unsupported architecture ' + args.architecture)
	# add output layers.
	model.add( layers.Dense(1) )
	model.add( layers.Activation('sigmoid') )
	
	optimizer = parse_optimizer(args)

	model.compile(optimizer=optimizer, loss=binary_crossentropy, metrics=['accuracy'])
	model.summary()
	"""
		Training
	"""
	history = History()
	try:
		model.fit(X, Y, batch_size=BATCH_SIZE, shuffle=True, \
			validation_split=0.1, verbose=1, nb_epoch=args.epochs, \
			callbacks=[ history, \
			EarlyStopping(min_delta=0.0005, patience=2, verbose=True)])

	except KeyboardInterrupt as e:
		print('Training Stopped Early')

	"""
		Evaluations
	"""
	print(args)
	print(params)
	Y_train_preds = model.predict(X, verbose=False)
	Y_val_preds = model.predict(X_val, verbose=False)

	PREFIX_FILENAME = args.outfile + args.embedding + args.architecture

	acc_train = evaluate.evaluate(Y, Y_train_preds, title='Training')
	acc_val = evaluate.evaluate(Y_val, Y_val_preds, title='Validation')
	f1_, _ = visualization.training_plot(history, args.outfile, metric='loss')
	f2_, _ = visualization.training_plot(history, args.outfile, metric='acc')
	
	json.dump(history.history, open(PREFIX_FILENAME + '_history.json', 'w'))

	fpr_train, tpr_train, auc_train = evaluate.roc_auc(Y, Y_train_preds)
	fpr_val, tpr_val, auc_val = evaluate.roc_auc(Y_val, Y_val_preds)

	f3_, _ = visualization.plot_auc([fpr_train, fpr_val], \
									[tpr_train, tpr_val], \
									[auc_train, auc_val], \
									labels=['Training Set', 'Validation Set'])

	# plt.show()
	f1_.savefig(PREFIX_FILENAME + '_loss.pdf')
	f2_.savefig(PREFIX_FILENAME + '_acc.pdf')
	f3_.savefig(PREFIX_FILENAME + '_auc.pdf')
	model.save(PREFIX_FILENAME + '_model.hdf5')

	if args.use_test:
		o = {
			'fpr': fpr_val.tolist(),
			'tpr': tpr_val.tolist(),
			'auc': auc_val.tolist()
		}

		json.dump(o, open(PREFIX_FILENAME+'.auc.json', 'w'))

		log_result(acc_train, acc_val, auc_val.tolist(), tpr_val.tolist(), fpr_val.tolist(), args)

if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	named_args = parser.add_argument_group('named arguments')

	named_args.add_argument('-e', '--embedding', metavar='|',
						help="""Representation of the data
						one of: 'glove', 'learned', 'tfidf', 'counts'""",
						required=True)

	named_args.add_argument('-o', '--outfile', metavar='|',
						help='File to store trained model',
						required=False, type=str, default='./training_results')

	named_args.add_argument('-n', '--epochs', metavar='|',
						help='Number of epochs to train network',
						required=True, type=int)

	named_args.add_argument('-a', '--architecture', metavar='|',
						help='Network Architecture to use (FNN, CNN, LSTM)',
						required=True)

	named_args.add_argument('-p', '--params', metavar='|',
						help='Parameter file to use to build the architecture',
						required=True)

	named_args.add_argument('-t', '--optimizer', metavar='|',
						help='Optimizer to use',
						required=False, default='adam')

	named_args.add_argument('-lr', '--learningrate', metavar='|',
						help='Optimizer to use',
						required=False, default=0.01, type=float)
	named_args.add_argument('-ut', '--use-test', metavar='|',
						help='Use the test set to evaluate the model',
						required=False, default=False, type=bool)
	named_args = parser.add_argument_group('embedding arguments')
	named_args.add_argument('-v', '--vocabsize', metavar='|',
						help="""Number of words in vocabulary""",
						required=False, type=int, default=5000)

	named_args.add_argument('-es', '--embeddingsize', metavar='|',
						help="""Size of embedding dimension""",
						required=False, type=int, default=50)

	named_args.add_argument('-nmi', '--ngram-min', metavar='|',
						help="""Minmum n-gram size""",
						required=False, type=int, default=1)

	named_args.add_argument('-nmx', '--ngram-max', metavar='|',
						help="""Maximum ngram size""",
						required=False, type=int, default=1)

	named_args.add_argument('-et', '--embedding-trainable', metavar='|',
						help="""Defines if the embedding is trainable""",
						required=False, type=bool, default=True)

	named_args.add_argument('-ss', '--subset', metavar='|',
						help="""Uses a subset of data to do training""",
						required=False, type=int, default=20000)

	args = parser.parse_args()
	if args.use_test:
		print('Testing mode')
	main(args)
