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

featureNum = 4096
simulation_time = 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, 'data/month_num=12/')
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 testi in range(startmark, endmark):
    print(testi, end=" ")
    start_time = time.time()
    x_train = np.loadtxt(os.path.join(default_path, 'train_features/train_features(' + str(testi) + ').txt'))
    z_train = np.loadtxt(os.path.join(default_path, 'train_demands/train_demands(' + str(testi) + ').txt'))


    regr_weight = Ridge(alpha=1.0)
    regr_weight.fit(x_train, z_train)

    x_test = np.loadtxt(os.path.join(default_path, 'test_features/test_features(' + str(testi) + ').txt'))
    z_predict = regr_weight.predict(x_test)
    z_real = np.loadtxt(os.path.join(default_path, 'test_demands/test_demands(' + str(testi) + ').txt'))
    res_weight = np.vstack((z_real, z_predict))
    res_weight = res_weight.T
    np.savetxt(os.path.join(default_path, 'Ridge/Ridge_demands(' + str(testi) + ').txt'), res_weight, fmt="%.2f")
    end_time = time.time()
    runtime = end_time - start_time
    print("MSE: ", mean_squared_error(z_real, z_predict), end=" ")
    print("training time: ", runtime)


print("knn")
from sklearn.neighbors import KNeighborsRegressor
k = 5
for testi in range(startmark, endmark):
    print(testi, end=" ")
    start_time = time.time()
    x_train = np.loadtxt(os.path.join(default_path, 'train_features/train_features(' + str(testi) + ').txt'))
    z_train = np.loadtxt(os.path.join(default_path, 'train_demands/train_demands(' + str(testi) + ').txt'))

    knn_weight = KNeighborsRegressor(k)
    knn_weight.fit(x_train, z_train)

    x_test = np.loadtxt(os.path.join(default_path, 'test_features/test_features(' + str(testi) + ').txt'))

    z_predict = knn_weight.predict(x_test)
    z_real = np.loadtxt(os.path.join(default_path, 'test_demands/test_demands(' + str(testi) + ').txt'))
    res_weight = np.vstack((z_real, z_predict))
    res_weight = res_weight.T
    np.savetxt(os.path.join(default_path, 'knn/knn_demands(' + str(testi) + ').txt'), res_weight, fmt="%.2f")
    end_time = time.time()
    runtime = end_time - start_time
    print("MSE: ", mean_squared_error(z_real, z_predict), end=" ")
    print("training time: ", runtime)



print("CART")
from sklearn import tree
for testi in range(startmark, endmark):
    print(testi, end=" ")
    start_time = time.time()
    x_train = np.loadtxt(os.path.join(default_path, 'train_features/train_features(' + str(testi) + ').txt'))
    z_train = np.loadtxt(os.path.join(default_path, 'train_demands/train_demands(' + str(testi) + ').txt'))

    tree_reg_weight = tree.DecisionTreeRegressor()
    tree_reg_weight.fit(x_train, z_train)

    x_test = np.loadtxt(os.path.join(default_path, 'test_features/test_features(' + str(testi) + ').txt'))
    z_predict = tree_reg_weight.predict(x_test)
    z_real = np.loadtxt(os.path.join(default_path, 'test_demands/test_demands(' + str(testi) + ').txt'))
    res_weight = np.vstack((z_real, z_predict))
    res_weight = res_weight.T
    np.savetxt(os.path.join(default_path, 'CART/CART_demands(' + str(testi) + ').txt'), res_weight, fmt="%.2f")
    end_time = time.time()
    runtime = end_time - start_time
    print("MSE: ", mean_squared_error(z_real, z_predict), end=" ")
    print("training time: ", runtime)



print("RF")
from sklearn import ensemble
for testi in range(startmark, endmark):
    print(testi, end=" ")
    start_time = time.time()
    x_train = np.loadtxt(os.path.join(default_path, 'train_features/train_features(' + str(testi) + ').txt'))
    z_train = np.loadtxt(os.path.join(default_path, 'train_demands/train_demands(' + str(testi) + ').txt'))

    rf_reg_weight = ensemble.RandomForestRegressor(n_estimators=100)
    rf_reg_weight.fit(x_train, z_train)

    x_test = np.loadtxt(os.path.join(default_path, 'test_features/test_features(' + str(testi) + ').txt'))
    z_predict = rf_reg_weight.predict(x_test)
    z_real = np.loadtxt(os.path.join(default_path, 'test_demands/test_demands(' + str(testi) + ').txt'))
    res_weight = np.vstack((z_real, z_predict))
    res_weight = res_weight.T
    np.savetxt(os.path.join(default_path, 'RF/RF_demands(' + str(testi) + ').txt'), res_weight, fmt="%.2f")
    end_time = time.time()
    runtime = end_time - start_time
    print("MSE: ", mean_squared_error(z_real, z_predict), end=" ")
    print("training time: ", runtime)
