import lzma

from sklearn.feature_extraction.text import TfidfVectorizer
# from sklearn.metrics import accuracy_score
from sklearn.linear_model import LogisticRegression
from stop_words import get_stop_words


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(line)
    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 classify_data(train):
    x_data = get_data(f"{train}/in.tsv", "xz")
    Y_data = get_data(f"{train}/expected.tsv", "tsv")

    custom_stop_words = get_stop_words("pl")
    vectorizer = TfidfVectorizer(stop_words=custom_stop_words)
    X_data = vectorizer.fit_transform(x_data)

    logreg = LogisticRegression(max_iter=1000)
    y_pred = logreg.fit(X_data, Y_data)

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

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


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

print(accuracy_score(y_true, y_pred))"""