from sklearn.linear_model import LinearRegression
from sklearn.feature_extraction.text import TfidfVectorizer
import pandas as pd

def regresja(path):

    # dnae do modelu
    train = pd.read_csv('train/train.tsv', sep='\t', header=None)
    y_train = train[0]
    x_train = train[4]

    # dane do predykcji
    x_predict = []
    with open(f"{path}/in.tsv", encoding='utf-8') as f:
        for line in f:
            x_predict.append(line)

    # tfidf
    vectorizer = TfidfVectorizer()
    x_train = vectorizer.fit_transform(x_train)

    x_predict = vectorizer.transform(x_predict)

    # model regresji
    model = LinearRegression()
    model.fit(x_train, y_train)

    # przewidywanie wyniku
    y_predict = model.predict(x_predict)
    pd.DataFrame(y_predict).to_csv(f"{path}/out.tsv", header=False, index=None)


regresja("dev-0")
regresja("dev-1")
regresja("test-A")