#!/usr/bin/env python

import os
import math
import pandas as pd
from sklearn.metrics import mean_squared_error
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import make_pipeline
from stop_words import get_stop_words


def rmse(y_expected, y_predicted):
    return math.sqrt(mean_squared_error(y_expected, y_predicted))


def read_tsv(path, filename, col_names, cutoff=None):
    i = 0
    data = []
    with open(os.path.join(path, filename), 'r', encoding='utf-8') as file:
        for line in file:
            if cutoff and i >= cutoff:
                break
            data.append(line.strip().split('\t'))
            i += 1

    df = pd.DataFrame(data, columns=col_names)
    return df


def evaluate_and_save(path, file, model):
    df = read_tsv(path, file, ['y'])['y'].values
    y_predicted = model.predict(df)
    # y_expected = read_tsv(path, 'expected.tsv', ['y'])['y'].values
    # print('score: ', rmse(y_expected, y_predicted))
    with open(os.path.join(path, 'out.tsv'), 'w') as f:
        for value in y_predicted:
            f.write(f'{value}\n')


def main():
    train_data = read_tsv('train', 'train.tsv', ['date_start', 'date_end', 'x1', 'x2', 'text'], 50000)
    train_data['date'] = (train_data['date_start'].astype(float) + train_data['date_end'].astype(float)) / 2
    train_x = train_data['text'].values
    train_y = train_data['date'].values
    vectorizer = TfidfVectorizer(stop_words=get_stop_words('polish'))
    model = LinearRegression()
    pipe = make_pipeline(vectorizer, model)
    pipe.fit(train_x, train_y)
    # print('training done')
    for path, file in [('dev-0', 'in.tsv'), ('dev-1', 'in.tsv'), ('test-A', 'in.tsv')]:
        evaluate_and_save(path, file, pipe)


if __name__ == '__main__':
    main()