from __future__ import print_function, division
import nltk
import os
import random
from collections import Counter
from nltk import word_tokenize, WordNetLemmatizer
from nltk.corpus import stopwords
from nltk import NaiveBayesClassifier, classify
from sklearn.svm import SVC
import numpy as np
import pickle

stoplist = stopwords.words('english')
#stoplist.append('gunay dhorum na')

def init_lists(folder):
    a_list = []
    file_list = os.listdir(folder)
    for a_file in file_list:
        f = open(folder + a_file, 'r')
        a_list.append(f.read())
    f.close()
    return a_list

def preprocess(sentence):
    lemmatizer = WordNetLemmatizer()
    return [lemmatizer.lemmatize(word.lower()) for word in word_tokenize(unicode(sentence, errors='ignore'))]

def get_dictionary(all_emails):
	emails = [a[0] for a in all_emails]
	all_words = []
	for mail in emails:
		all_words += mail.split()
	dictionary = Counter(all_words)

	list_to_remove = dictionary.keys()
	for item in list_to_remove:
		if item.isalpha() == False:
			del dictionary[item]
		elif len(item) == 1:
			del dictionary[item]

	dictionary = dictionary.most_common(2500)
	return dictionary

def extract_feature(dictionary, words):
	temp = words.split()
	dictionary_words = [a[0] for a in dictionary]
	features = [words.count(word) if word in temp else 0 for word in dictionary_words]
	return features

def train(features, samples_proportion):
    train_size = int(len(features) * samples_proportion)
    train_set, test_set = features[:train_size], features[train_size:]
    x_train, y_train = [a[0] for a in train_set], [b[1] for b in train_set]
    classifier = SVC(kernel="rbf",C=120, verbose=True)
    classifier.fit(x_train, y_train)
    pickle.dump(classifier,open('trained_data.sav','wb'))
    print('Model has been saved with filename trained_data.sav')
    return train_set, test_set, classifier


spam = init_lists('data/neww/')
ham = init_lists('data/new/')
all_emails = [(email, 1) for email in spam]
all_emails += [(email, 0) for email in ham]
random.shuffle(all_emails)
dictionary = get_dictionary(all_emails)
all_features = [(extract_feature(dictionary, email), label) for (email, label) in all_emails]
train_set, test_set, classifier = train(all_features, 1)
