import os
import sys
import numpy as np
from sklearn import datasets
from sklearn.metrics import mean_squared_error
import math, time
import linear_relax as LP_relax_file
from sklearn import linear_model
from sklearn.linear_model import Ridge

featureNum = 8
testNum = 10

startmark = int(sys.argv[1])
endmark = int(sys.argv[2])

dataset_path = os.path.abspath(os.path.dirname(os.getcwd()))
default_path = os.path.join(dataset_path, 'day_num=7/')
LP_relax_file.mkdir(default_path, 'Ridge')
LP_relax_file.mkdir(default_path, 'knn')
LP_relax_file.mkdir(default_path, 'CART')
LP_relax_file.mkdir(default_path, 'RF')

print("Ridge")
for n in range(startmark, endmark):
    start_time = time.time()
    trainData = np.loadtxt(os.path.join(default_path, 'train/train(' + str(n) + ').txt'))
    x_train = trainData[:, 1:featureNum+1]
    y_train = trainData[:, featureNum+1]

    regr = Ridge(alpha=1.0)
    regr.fit(x_train, y_train)

    testData = np.loadtxt(os.path.join(default_path, 'test/test(' + str(n) + ').txt'))
    benchId = testData[:, 0]
    x_test = testData[:, 1:featureNum+1]
    y_real = testData[:, featureNum+1]
    y_predict = regr.predict(x_test)
    res = np.vstack((benchId, y_real, y_predict))
    res = res.T
    end_time = time.time()
    np.savetxt(os.path.join(default_path, 'Ridge/Ridge(' + str(n) + ').txt'), res, fmt="%.2f")
    print("MSE: ", mean_squared_error(y_real, y_predict), "runtime: ", end_time - start_time)


print("k-NN")
from sklearn.neighbors import KNeighborsRegressor
k = 5
for n in range(startmark, endmark):
    start_time = time.time()
    trainData = np.loadtxt(os.path.join(default_path, 'train/train(' + str(n) + ').txt'))
    x_train = trainData[:, 1:featureNum+1]
    y_train = trainData[:, featureNum+1]

    knn = KNeighborsRegressor(k)
    knn.fit(x_train, y_train)

    testData = np.loadtxt(os.path.join(default_path, 'test/test(' + str(n) + ').txt'))
    benchId = testData[:, 0]
    x_test = testData[:, 1:featureNum+1]
    y_real = testData[:, featureNum+1]

    y_predict = knn.predict(x_test)
    res = np.vstack((benchId, y_real, y_predict))
    res = res.T
    end_time = time.time()
    np.savetxt(os.path.join(default_path, 'knn/knn(' + str(n) + ').txt'), res, fmt="%.2f")
    print("MSE: ", mean_squared_error(y_real, y_predict), "runtime: ", end_time - start_time)


print("CART")
from sklearn import tree
for n in range(startmark, endmark):
    start_time = time.time()
    trainData = np.loadtxt(os.path.join(default_path, 'train/train(' + str(n) + ').txt'))
    x_train = trainData[:, 1:featureNum+1]
    y_train = trainData[:, featureNum+1]

    tree_reg = tree.DecisionTreeRegressor()
    tree_reg.fit(x_train, y_train)

    testData = np.loadtxt(os.path.join(default_path, 'test/test(' + str(n) + ').txt'))
    benchId = testData[:, 0]
    x_test = testData[:, 1:featureNum+1]
    y_real = testData[:, featureNum+1]

    y_predict = tree_reg.predict(x_test)
    res = np.vstack((benchId, y_real, y_predict))
    res = res.T
    end_time = time.time()
    np.savetxt(os.path.join(default_path, 'CART/CART(' + str(n) + ').txt'), res, fmt="%.2f")
    print("MSE: ", mean_squared_error(y_real, y_predict), "runtime: ", end_time - start_time)


print("RF")
from sklearn import ensemble
for n in range(startmark, endmark):
    start_time = time.time()
    trainData = np.loadtxt(os.path.join(default_path, 'train/train(' + str(n) + ').txt'))
    x_train = trainData[:, 1:featureNum+1]
    y_train = trainData[:, featureNum+1]

    rf_reg = ensemble.RandomForestRegressor(n_estimators=100)
    rf_reg.fit(x_train, y_train)

    testData = np.loadtxt(os.path.join(default_path, 'test/test(' + str(n) + ').txt'))
    benchId = testData[:, 0]
    x_test = testData[:, 1:featureNum+1]
    y_real = testData[:, featureNum+1]
    y_predict = rf_reg.predict(x_test)
    res = np.vstack((benchId, y_real, y_predict))
    res = res.T
    end_time = time.time()
    np.savetxt(os.path.join(default_path, 'RF/RF(' + str(n) + ').txt'), res, fmt="%.2f")
    print("MSE: ", mean_squared_error(y_real, y_predict), "runtime: ", end_time - start_time)
