'''
Multinomial Naive Bayes Classifier
==================================
'''

from BagOfWords import BagOfWords

import csv

import gensim
from gensim.models.doc2vec import Doc2Vec, TaggedDocument
import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_selection import SelectPercentile
from sklearn.metrics import recall_score, precision_score
import sklearn
from sklearn.model_selection import StratifiedKFold
from sklearn.naive_bayes import MultinomialNB

class MultinomialNaiveBayes:

	def make_mnb(dataset, sklearn_cv=True, percentile=100):
		'''fits naive bayes model with StratifiedKFold
		'''
		vector_size=150

		def read_corpus(data, tokens_only=False):
			list_of_lists = []
			for i, text in enumerate(data):
				if tokens_only:
					list_of_lists.append(BagOfWords.extract_words(text))
				else:
					# For training data, add tags
					list_of_lists.append(gensim.models.doc2vec.TaggedDocument(BagOfWords.extract_words(text), [i]))
			return list_of_lists

		print('# starting multinomial naive bayes')
		print('# ...')

		# split data into text and label set
		# join title and text
		X = dataset['Title'] + '. ' + dataset['Text']
		y = dataset['Label']

		# use stratified k-fold cross-validation as split method
		skf = StratifiedKFold(n_splits = 10, shuffle=True, random_state=5)

		classifier = MultinomialNB(alpha=1.0e-10,
								   fit_prior=False,
								   class_prior=None)

		# metrics
		recall_scores = []
		precision_scores = []
		f1_scores = []

		# probabilities of each class (of each fold)
		#class_prob = []
		# counts number of training samples observed in each class 
		#class_counts = []

		# for each fold
		n = 0
		for train, test in skf.split(X,y):

			n += 1
			print('# split no. ' + str(n))

			# train model with gensim	   
			training_data = read_corpus(X[train], tokens_only=False)
			testing_data = read_corpus(X[test], tokens_only=True)
			all_data = read_corpus(X, tokens_only=False)

			# instantiate a Doc2Vec object
			doc2vec_model = Doc2Vec(training_data, vector_size=5, window=2, min_count=1, workers=4)

			print(doc2vec_model.docvecs[0])
			print(doc2vec_model.docvecs[1])
			print(doc2vec_model.docvecs[2])
			
			training_data = [doc2vec_model.docvecs[i] for i in range(len(training_data))]
			testing_data = [doc2vec_model.infer_vector(vector) for vector in testing_data]
			
			#fit classifier
			classifier.fit(training_data, y[train])
			#predict class
			predictions_train = classifier.predict(training_data)
			predictions_test = classifier.predict(testing_data)

			#print and store metrics
			rec = recall_score(y[test], predictions_test, average='weighted')
			print('rec: ' + str(rec))
			recall_scores.append(rec)
			prec = precision_score(y[test], predictions_test, average='weighted')
			print('prec: ' + str(prec))
			print('#')
			precision_scores.append(prec)
			# equation for f1 score
			f1_scores.append(2 * (prec * rec)/(prec + rec))

		##########################
		# probability estimates for the test vector (testing_data)
		class_probs = classifier.predict_proba(testing_data)

		# number of samples encountered for each class during fitting
		# this value is weighted by the sample weight when provided
		class_count = classifier.class_count_

		# classes in order used
		classes = classifier.classes_
		
		print('average: recall, precision, f1 score')
		print(sum(recall_scores)/10, sum(precision_scores)/10, sum(f1_scores)/10)
		
		
		# return classes and vector of class estimates
		return recall_scores, precision_scores, f1_scores, class_probs

if __name__ == '__main__':

	# read csv file
	print('# reading dataset')
	print('# ...')

	# read current data set from csv
	df = pd.read_csv('../data/interactive_labeling_round_11.csv',
			  sep='|',
			  usecols=range(1,13), # drop first column 'unnamed'
			  encoding='utf-8',
			  quoting=csv.QUOTE_NONNUMERIC,
			  quotechar='\'')

	# select only labeled articles
	MultinomialNaiveBayes.make_mnb(df.loc[df['Label'] != -1][:100].reset_index(drop=True), sklearn_cv=False, percentile=100)