import lzma

from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import TfidfVectorizer
# from sklearn.metrics import accuracy_score


def get_data(file_name, data_type):
    lines = []
    if data_type == "tsv":
        with open(file_name, encoding="utf-8") as file:
            for line in file.readlines():
                lines.append(int(line.replace("\n", "")))
    else:
        with lzma.open(f"{file_name}.{data_type}") as file:
            for line in file.readlines():
                lines.append(line.rstrip().decode("utf-8"))
    return lines


def bayes(train):
    x_data = get_data(f"{train}/in.tsv", "xz")
    Y_data = get_data(f"{train}/expected.tsv", "tsv")

    vectorizer = TfidfVectorizer(stop_words="english")
    X_data = vectorizer.fit_transform(x_data)

    clf = MultinomialNB()
    y_pred = clf.fit(X_data, Y_data)

    for predct in ["test-A", "dev-0"]:
        Y_test = get_data(f"{predct}/in.tsv", "xz")
        y_prediction = y_pred.predict(vectorizer.transform(Y_test))

        with open(f"{predct}\out.tsv", "w", encoding="UTF-8") as file_out:
            for single_pred in y_prediction:
                file_out.writelines(f"{str(single_pred)}\n")


bayes("train")

'''y_true = []
with open("dev-0/expected.tsv", encoding='utf-8') as file:
            for line in file.readlines():
                y_true.append(line)
y_pred = []
with open("dev-0/out.tsv", encoding='utf-8') as file:
            for line in file.readlines():
                y_pred.append(line)
print(accuracy_score(y_true, y_pred))'''