#!  /usr/bin/python3

import sklearn.datasets
import sklearn.metrics
import sklearn.cross_validation
import sklearn.svm
import sklearn.naive_bayes
import sklearn.neighbors
import scipy.io
import scipy.sparse
from colorama import init
from termcolor import colored
import sys
import os
import glob

dir_path = "../small_dataset"
files = sklearn.datasets.load_files(dir_path,shuffle=False)

count_vector = sklearn.feature_extraction.text.CountVectorizer()
word_counts = count_vector.fit_transform(files.data)
x = count_vector.vocabulary_

voc = {x[k]:k for k in x}
f = open("vocabulary.dict","w")
for k in voc:
	f.write("{0}:{1}\n".format(k,voc[k]))
f.close()

f = open("filename.txt","w")
for filename in files.filenames:
	f.write("{0}\n".format(filename))
f.close()

#import operator
#sorted_x = sorted(x.items(), key=operator.itemgetter(1))
#print(sorted_x)

tf_transformer = sklearn.feature_extraction.text.TfidfTransformer(use_idf=True).fit(word_counts)
X = tf_transformer.transform(word_counts)

cx = scipy.sparse.coo_matrix(X)

for i,j,v in zip(cx.row, cx.col, cx.data):
	if i == 0:
	    print("(%s, %d), %s" % (voc[j],j,v))

#scipy.io.mmwrite("data.matrix",X)
clf = sklearn.svm.LinearSVC()

X_train, X_test, y_train, y_test = sklearn.cross_validation.train_test_split(X, files.target, test_size = 0.2)

clf.fit(X_train, y_train)
y_predicted = clf.predict(X_test)

print(sklearn.metrics.classification_report(y_test, y_predicted, target_names=files.target_names))
