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

def getData(path):
    with open(path, encoding="utf-8") as source:
          return source.readlines()

def changeToDf(input):
  return pd.read_csv(input, sep="\t")

vectorizer = TfidfVectorizer()
linear = LinearRegression()
train = pd.read_csv("./train/train.tsv", sep="\t", names=['start_date', 'end_date', 'title', 'sort_title', 'data'])
mean = (train['start_date'] + train['end_date']) / 2
tv = vectorizer.fit_transform(train['data'])
linear.fit(tv, mean)

def calculateResult(in_, out):
  tmp = getData(in_)
  df = pd.DataFrame(data = tmp)
  data = vectorizer.transform(df[0])
  evaluate = linear.predict(data)
  np.savetxt(out, evaluate, fmt='%f', delimiter='\n')

calculateResult("./dev-0/in.tsv", "./dev-0/out.tsv")
calculateResult("./dev-1/in.tsv", "./dev-1/out.tsv")
calculateResult("./test-A/in.tsv", "./test-A/out.tsv")