import gzip
import io
import pandas as pd 
import numpy as np 
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline
from sklearn.metrics import accuracy_score

def read_data_gz(baseUrl):
    f = gzip.open(baseUrl,'r')
    data_unzip = f.read()
    data = pd.read_table(io.StringIO(data_unzip.decode('utf-8')), error_bad_lines=False, header= None)
    return data

baseUrl = '/home/przemek/ekstrakcja/sport-text-classification-ball-ISI-public/'
data = read_data_gz(baseUrl + 'train/train.tsv.gz')

y_train = data[0].values
x_train = data[1].values

model = make_pipeline(TfidfVectorizer(), MultinomialNB())
model.fit(x_train, y_train)

# dev-0
x_dev = pd.read_table(baseUrl + 'dev-0/in.tsv', error_bad_lines=False, header= None)
x_dev = x_dev[0].values

y_pred = model.predict(x_dev)
y_pred.tofile(baseUrl + 'dev-0/out.tsv', sep='\n')
# --------------

# test-A
x_testA =  pd.read_table(baseUrl + '/test-A/in.tsv', error_bad_lines=False, header= None)
x_testA= x_testA[0].values

y_predA = model.predict(x_testA)
y_predA.tofile(baseUrl + 'test-A/out.tsv', sep='\n')
# --------------