import pandas as pd
import numpy as np
import os, csv, time
import spacy # has pos tagger, word vector stuff and a bunch of other tools, very fast as well
from keras.preprocessing.sequence import pad_sequences
from sklearn.model_selection import train_test_split
import json
from tqdm import tqdm

### model
from keras.models import Sequential
from keras.layers import (Embedding, Activation, LSTM, Dense,Dropout)
from keras.utils import to_categorical

top_path = '/complete/'
path_to_data = os.getcwd() + '/' + top_path

global_timer = time.time()

files = os.listdir(path_to_data)
files = list(filter(lambda k: '.tsv' in k,files)) # keep only tsv files

temp_frames = []

df = pd.DataFrame()

for file in files:
	print("Reading {}".format(file))

	temp_df = pd.read_csv(path_to_data+file,delimiter='	',quoting=csv.QUOTE_NONE) 
	temp_frames.append(temp_df)

df = pd.concat(temp_frames,ignore_index=True)

# adding simpler labels
def labeling(row):
	if row['annotation'] == 'NoArgument':
		return 0
	elif row['annotation'] =='Argument_for':
		return 1
	elif row['annotation'] =='Argument_against':
		return 2

# apply labels
df['label'] = df.apply(lambda row: labeling(row),axis=1)

### Once we have the data, time to initialize the spacy parser
l_model = 'en_core_web_lg'

print("Loading model {}".format(l_model))

start_time = time.time()

nlp = spacy.load(l_model) # this is the large sized model that also supports word vectors (for some reason loads faster than medium size model)

print("Done. Loading took {} seconds".format((time.time()-start_time)))

vectors = []

raw_sentences = df['sentence'].tolist()
labels = df['label'].tolist()

processed_sentences = []
max_len = 100
num_over_limit = 0

print("Processing sentences using spaCy")

for s in tqdm(range(len(raw_sentences))): # len(raw_sentences) goes here normally
	processed = nlp(raw_sentences[s])
	processed_sentences.append(processed)

	temp_vectors = []

	for word in processed:
		if(not word.pos_ in ['PUNCT','SYM','X']): # not word.is_stop and 
			temp_vectors.append(word.lex_id)

	if(len(temp_vectors) > max_len):
		temp_vectors = temp_vectors[:max_len] # if sequence is longer, simply truncate it (should not be needed, but ok)
		num_over_limit += 1

	vectors.append(temp_vectors)
	
print("Number of sentences found to be over max length of {} is {}".format(max_len,num_over_limit))

# set backend
os.environ['KERAS_BACKEND']='tensorflow'

# some variables we're going to need for the model
vocab_len = len(nlp.vocab.vectors.data) # just in case
vector_len = 300
batch_size = 32

print("Maximum sequence length = {}".format(max_len))

# last bit of pre-processing
print("Padding vectors")

padded_vectors = pad_sequences(vectors,maxlen=max_len, padding='pre')

print("Splitting data")

labels = to_categorical(labels)
X_train, X_test, Y_train, Y_test = train_test_split(padded_vectors,labels,test_size=0.2,random_state=42) # try a split that may not overfit?

print("Size of training data: {}".format(len(X_train)))

# initialize model
print("Initializing Keras Model")

model = Sequential()
model.add(Embedding(vocab_len, 300, weights = [nlp.vocab.vectors.data], input_length = max_len, trainable = False))
model.add(LSTM(max_len, dropout = 0.2, recurrent_dropout = 0.2, use_bias = True))
model.add(Dense(3, activation = 'softmax'))
model.compile(loss = 'categorical_crossentropy', optimizer = 'adam', metrics = ['accuracy'])

print(model.summary())

# train the model
print("Training model")

model_timing = time.time()
model.fit(X_train, Y_train, batch_size = batch_size, epochs = 3, verbose = 2) # defalut epoch = 30, Changed for testing puropse now. 

print("Model Trained. Time elapsed: {}".format(time.time() - model_timing))

# score model
score,acc = model.evaluate(X_test, Y_test, verbose = 2, batch_size = batch_size)

print('Test accuracy: {}\nModel Score: {}'.format(acc,score))

# Save model to disk
def save_model(model):
	model_save = 'model_best'

	print("Converting model to json")

	model_json = model.to_json()

	print("Saving model to {}".format(model_save))

	with open(model_save+'.json','w+') as json_file:
		json_file.write(model_json)

	print("Saving model weights")

	model.save_weights(model_save+'.h5')

	print("Saved model to disk")

save_model(model)

print("Done.")
print("Total time elapsed: {}s".format((time.time()-global_timer)))


