import sys
import parsetags
import numpy as np
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.naive_bayes import MultinomialNB
from sklearn import svm
from sklearn.decomposition import PCA as PCA
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import nltk
from nltk.stem import LancasterStemmer
import re

def get_top_tags(segment_tags, n):
	tag_freqs = {}
	for tag_list in segment_tags.values():
		for tag in tag_list:
			if tag not in tag_freqs:
				tag_freqs[tag] = 0
			tag_freqs[tag] += 1
	return ['NULL'] + sorted(tag_freqs.keys(), key=lambda x: tag_freqs[x])[-n:]

def get_common_words(n=100):
	try:
		file_content = open(sys.argv[3]).read()
		common_words = nltk.word_tokenize(file_content)
	except IndexError:
		return None
	return set(common_words[:n])

def get_named_entities():
	try:
		file_content = open(sys.argv[2]).read()
		named_entities = nltk.word_tokenize(file_content)
	except IndexError:
		return None
	return set(named_entities)

def filter_segments(segment_tags, ntags):
	filtered_segtags = {}
	for segment in segment_tags:
		# np.random.shuffle(segment_tags[segment])
		for tag in segment_tags[segment]:
			if tag not in ntags: continue
			filtered_segtags[segment] = ntags.index(tag)
		if segment not in filtered_segtags:
			filtered_segtags[segment] = 0
	return filtered_segtags

def increase_num_segments(segment_tags, n, length=1000):
	new_segment_tags = {}
	segments = sorted(segment_tags.keys(), key=len)
	lengths = np.array([len(seg) for seg in segments])
	dist = lengths/np.sum(lengths)
	random_segments = np.random.choice(segments, size=n, p=dist)
	for segment in random_segments:
		new_segment = segment
		if len(new_segment) > length:
			index = np.random.randint(0, len(new_segment)-length)
			new_segment = new_segment[index:index+length]
		new_segment_tags[new_segment] = segment_tags[segment]
	return new_segment_tags

def named_entity_reduction(segment_tags, named_entities, common_words):
	punctuation = [',', '.', "'", '?', ';', ':', '!', '(', ')', '`', '--'
			   '\xe2', '\x80', '\x94', '\x99']
	new_segments = []
	segments = list(segment_tags.keys())
	for segment in segments:
		new_segment = ''
		tokens = nltk.word_tokenize(segment)
		for token in tokens:
			if token in punctuation: continue
			if token.lower() in common_words: continue
			if token not in named_entities: continue
			new_segment += token + ' '
		new_segments.append(new_segment)
	new_segment_tags = {}
	for i in range(len(segments)):
		new_segment_tags[new_segments[i]] = segment_tags[segments[i]]
	return new_segment_tags

def stemming_reduction(segment_tags):
	punctuation = [',', '.', "'", '?', ';', ':', '!', '(', ')', '`', '--'
					'\xe2', '\x80', '\x94', '\x99']
	new_segments = []
	stemmer = LancasterStemmer()
	segments = list(segment_tags.keys())
	for segment in segments:
		new_segment = ''
		segment = re.sub(r'[^\x00-\x7f]',r'', segment)
		tokens = nltk.word_tokenize(segment)
		for token in tokens:
			if token in punctuation: continue
			try: 
				new_segment += stemmer.stem(token)+' '
			except UnicodeDecodeError:
				new_segment += ''
		new_segments.append(new_segment)
	stemmed_segment_tags = {}
	for i in range(len(segments)):
		stemmed_segment_tags[new_segments[i]] = segment_tags[segments[i]]
	return stemmed_segment_tags


def separate_segments(segment_tags, k):
	train = {}
	for segment in segment_tags.keys():
		if np.random.random() < k:
			train[segment] = segment_tags.pop(segment)
	return train, segment_tags


def bag_of_words(segment_tags, tfidf=False):
	#create matrix of word frequencies 
	segments = list(segment_tags.keys())
	vec = CountVectorizer()
	word_freqs = vec.fit_transform(segments).toarray()
	if tfidf:
		tfidf_transformer = TfidfTransformer()
		word_freqs = tfidf_transformer.fit_transform(word_freqs)
	labels = np.empty(shape=len(segments))
	for i in range(len(segments)):
		labels[i] = segment_tags[segments[i]]
	return word_freqs, labels, segments

def entity_bow(segment_tags, named_entities, common_words):
	punctuation = [',', '.', "'", '?', ';', ':', '!', '(', ')', '`', '--'
				   '\xe2', '\x80', '\x94', '\x99']
	new_segments = []
	segments = list(segment_tags.keys())
	for segment in segments:
		new_segment = ''
		tokens = nltk.word_tokenize(segment)
		for token in tokens:
			if token in punctuation: continue
			if token.lower() in common_words: continue
			if token not in named_entities: continue
			new_segment += token + ' '
		new_segments.append(new_segment)

	vec = CountVectorizer()
	word_freqs = vec.fit_transform(new_segments).toarray()
	tfidf_transformer = TfidfTransformer()
	X_train_tfidf = tfidf_transformer.fit_transform(word_freqs)
	print(word_freqs.shape, X_train_tfidf.shape)
	labels = np.empty(shape=len(segments))
	for i in range(len(segments)):
		labels[i] = segment_tags[segments[i]]
	return X_train_tfidf, labels, segments

def pca_plot(Xtrain, ytrain):
	#binary classification case
	X_reduced = Xtrain
	pca = PCA(3)
	X_pca = pca.fit_transform(X_reduced)
	ax = plt.axes(projection='3d')
	for i in range(X_pca.shape[0]):
		if ytrain[i] == 1:
			ax.scatter(X_pca[i, 0], X_pca[i, 2], X_pca[i, 1], 'o', color='blue')
		else:
			ax.scatter(X_pca[i, 0], X_pca[i, 2], X_pca[i,1], 'x', color='red')
	plt.show()

def randomize_order(X, y, segments):
	shuffled_segments = []
	indices = np.arange(len(segments))
	np.random.shuffle(indices)
	X, y = X[indices], y[indices]
	for i in indices:
		shuffled_segments.append(segments[i])
	return X, y, segments

def naive_bayes(segment_tags, k=0.5, normalize=False):
	X, y, segments = randomize_order(*bag_of_words(segment_tags, tfidf=normalize))
	num_examples = len(segments)
	Xtest, ytest = X[int(k*num_examples):, :], y[int(k*num_examples):]
	Xtrain, ytrain = X[:int(k*num_examples), :], y[:int(k*num_examples)]
	nb_classifier = MultinomialNB().fit(Xtrain, ytrain)
	nb_predicted_tags = nb_classifier.predict(Xtest)
	nb_success_rate = np.mean(nb_predicted_tags == ytest)
	return 1-nb_success_rate

def support_vector(segment_tags, k=0.5, normalize=False):
	X, y, segments = randomize_order(*bag_of_words(segment_tags, tfidf=normalize))
	num_examples = len(segments)
	Xtest, ytest = X[int(k*num_examples):, :], y[int(k*num_examples):]
	Xtrain, ytrain = X[:int(k*num_examples), :], y[:int(k*num_examples)]
	svm_classifier = svm.SVC()
	svm_classifier.fit(Xtrain, ytrain)
	svm_predicted_tags = svm_classifier.predict(Xtest)
	svm_success_rate = np.mean(svm_predicted_tags == ytest)
	return 1-svm_success_rate

if __name__ == "__main__":
	orig_segment_tags = parsetags.parse_tags(sys.argv[1])
	common_words = get_common_words()
	named_entities = get_named_entities()
	for sample_size in ['BIG']:
		if sample_size == 'BIG':
			segment_tags = increase_num_segments(orig_segment_tags, 3000, length=1000)
		orig_segment_tags = segment_tags
		for text_features in ['REGULAR', 'STEMMED', 'NAMED']:
			if text_features == 'STEMMED':
				segment_tags = stemming_reduction(orig_segment_tags)
			if text_features == 'NAMED':
				segment_tags = named_entity_reduction(orig_segment_tags, named_entities, common_words)
			for freq_feature in ['COUNT', 'TFIDF']:
				# ntags = get_top_tags(segment_tags, 7)
				print(sample_size, text_features, freq_feature)
				ntags = ['ETA', 'EHDRH', 'AFR']
				filtered_segtags = filter_segments(segment_tags, ntags)
				with open('Results/' + sample_size + '_' + text_features + '_' + freq_feature + '.txt', 'w') as f:
					for i in range(100):
						f.write(str(naive_bayes(filtered_segtags, normalize=(freq_feature is 'TFIDF'))) + '\n')
				
				
				


	# segment_tags = parsetags.parse_tags(sys.argv[1])
	# big_segment_tags = increase_num_segments(segment_tags, 3000, length=1000)
	# ntags = get_top_tags(segment_tags, 7)
	# for 
	# # ntags = ['NULL', 'ETA', 'EHDRH', 'AFR']
	
	
	# common_words = get_common_words()
	# named_entities = get_named_entities()
	# filtered_segtags = filter_segments(segment_tags, ntags)
	# #entity_bow(filtered_segtags, named_entities, common_words)
	# naive_bayes(filtered_segtags, named_entities, common_words, features=entity_bow)
	# naive_bayes(filtered_segtags)
	# support_vector(filtered_segtags)
	
	# predicted_tags = [ntags[int(np.round(nb_predicted_tags[i]))] for i in range(len(svm_predicted_tags))]
	# count = 0
	# print(ntags)
	# for i in range(len(predicted_tags)):
	# 	if predicted_tags[i] == 'NULL':
	# 		if all(tag not in segment_tags[shuffled_segments[i]] for tag in ntags):
	# 			count += 1
	# 	else:
	# 		if predicted_tags[i] in segment_tags[shuffled_segments[i]]:
	# 			count += 1
	# print(count/len(predicted_tags))



