#This is LSTM undersampled case

#python imports
import numpy as np
import random
import codecs
from nltk.corpus import stopwords
from os import listdir

#Keras imports
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import LSTM, Flatten
from keras.layers.embeddings import Embedding
from keras.preprocessing.text import one_hot
from keras.preprocessing.sequence import pad_sequences
from keras.preprocessing import sequence

#sklearn imports
from sklearn.metrics import precision_recall_fscore_support
from sklearn.metrics import f1_score
from sklearn.utils import class_weight


def try_lstm(X, Y):
	np.random.seed(7)
	top_words = 5000
	max_length = 100
	
	#Use one-hot encoding to encode the sentences
	encoded_sentences = [one_hot(d, top_words) for d in X]

	#Pad sentences if the length is greater than maxlen and truncate if length is lesser tan maxlen
    	padded_sentences = pad_sequences(encoded_sentences, maxlen = max_length, padding='post')

    	#Split into training and test sets
	X_train = padded_sentences[:3000]
	X_test = padded_sentences[3000:]
	y_train = Y[:3000]
	y_test = Y[3000:]

	#calculate class weights to for using them in loss function
	myclass_weight = class_weight.compute_class_weight('balanced', np.unique(y_train), y_train)
	class_weight_dict = dict(enumerate(myclass_weight))
	
	#Create the model
	#Use word embeddings of 32 dimensions, fewer dimensions in this case provide better performance due to smaller dataset
	embedding_vecor_length = 32
	model = Sequential()
	
	#Add the embedding layer at the beginning to accept input of max length
	model.add(Embedding(top_words, embedding_vecor_length, input_length=max_length))

	#Add three LSTM layers with 256, 128 and 16 units respectively and a dropout of 0.2.
	model.add(LSTM(256, dropout=0.2, recurrent_dropout=0.2, name='lstm1', return_sequences=True))
	model.add(LSTM(128, dropout=0.2, recurrent_dropout=0.2, name='lstm2', return_sequences=True))
	model.add(LSTM(16, dropout=0.2, recurrent_dropout=0.2, name='lstm3', return_sequences=True))
	model.add(Flatten())

	#Add 1 unit NN for classification
	model.add(Dense(1, activation='sigmoid'))

	#Train the model on the dataset
	model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
	print(model.summary())
	model.fit(X_train, y_train, epochs=1, batch_size=64, class_weight=class_weight_dict)

	# Final evaluation of the model
	scores = model.evaluate(X_test, y_test, verbose=0)

	#Predict the classes and the print Precision, Recall, Fscore and Accuracy
	a = model.predict_classes(X_test, verbose=1)
	print(precision_recall_fscore_support(y_test, a))
	print("F-score", f1_score(y_test, a))	
	print("Accuracy: %.2f%%" % (scores[1]*100))

#Take input folder and return dataset with labels
def generateData(dataFolder):
	stops = set(stopwords.words("english"))  
	X = []
	Y = []
	sentence = ''
	relevance = 0
	for fileName in listdir(dataFolder):
		f = codecs.open(dataFolder + fileName, encoding='utf-8')
		for line in f:
			if '\n' in line and len(line) == 2: #End of a Sentence
				if sentence != '':
					X.append(sentence)
					Y.append(relevance)
				#Reset sentence and relevance for next line
				sentence = ''
				relevance = 0 
			else:	#Same sentence still being continued
				token = line.split(' ')[0].lower()
				#Keep only words of letters and not stop words
				if token not in stops and token not in ",!.:<>":
					if sentence == '': 
						sentence = token
					else: 
						sentence += ' ' + token
				if line[:-1].split(' ')[-1] != 'O': #if it is annotated, then it is relevant 
					relevance = 1

	return X,Y

#Get the training data
X, Y = generateData("training_material/data/tokenized/")

#Find number of positives and negatives in the dataset and make it balanced
pos = []
neg = []
i = 0
for i in range(0, len(X)):
	if Y[i] == 1:
		pos.append((X[i], Y[i]))
	else:
		neg.append((X[i], Y[i]))
print len(pos)
print len(neg)

data = random.sample(pos, 2000)
data = data + random.sample(neg, 2000)
random.shuffle(data)

train = data[:3000]
test = data[3000:]
X_train = []
Y_train = []
X_test = []
Y_test = []
for each in train:
	X_train.append(each[0])
	Y_train.append(each[1])


for each in test:
	X_test.append(each[0])
	Y_test.append(each[1])


try_lstm(X_train + X_test, Y_train + Y_test)
