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

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/item_num=5, month_num=4/')
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')

featureNum = 4096
simulation_time = 10

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'))
    y_train = np.loadtxt(os.path.join(default_path, 'rescale_train_prices/rescale_train_prices(' + str(testi) + ').txt'))
    z_train = np.loadtxt(os.path.join(default_path, 'train_weights/train_weights(' + str(testi) + ').txt'))

    regr_value = Ridge(alpha=1.0)
    regr_value.fit(x_train, y_train)
    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'))
#    x_test = np.loadtxt('./CombOptNet/test_featureAll.txt')
    y_predict = regr_value.predict(x_test)
    y_real = np.loadtxt(os.path.join(default_path, 'rescale_test_prices/rescale_test_prices(' + str(testi) + ').txt'))
#    y_real = np.loadtxt('./CombOptNet/test_weightsAll.txt')
    res_value = np.vstack((y_real, y_predict))
    res_value = res_value.T
    np.savetxt(os.path.join(default_path, 'Ridge/Ridge_prices(' + str(testi) + ').txt'), res_value, fmt="%.2f")
    z_predict = regr_weight.predict(x_test)
    z_real = np.loadtxt(os.path.join(default_path, 'test_weights/test_weights(' + 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_weights(' + str(testi) + ').txt'), res_weight, fmt="%.2f")
    end_time = time.time()
    runtime = end_time - start_time
    print("training time: ", runtime)
    

print("k-NN")
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'))
    y_train = np.loadtxt(os.path.join(default_path, 'rescale_train_prices/rescale_train_prices(' + str(testi) + ').txt'))
    z_train = np.loadtxt(os.path.join(default_path, 'train_weights/train_weights(' + str(testi) + ').txt'))

    knn_value = KNeighborsRegressor(k)
    knn_value.fit(x_train, y_train)
    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'))
#    x_test = np.loadtxt('./CombOptNet/test_featureAll.txt')
    y_predict = knn_value.predict(x_test)
    y_real = np.loadtxt(os.path.join(default_path, 'rescale_test_prices/rescale_test_prices(' + str(testi) + ').txt'))
#    y_real = np.loadtxt('./CombOptNet/test_weightsAll.txt')
    res_value = np.vstack((y_real, y_predict))
    res_value = res_value.T
    np.savetxt(os.path.join(default_path, 'knn/knn_prices(' + str(testi) + ').txt'), res_value, fmt="%.2f")
    z_predict = knn_weight.predict(x_test)
    z_real = np.loadtxt(os.path.join(default_path, 'test_weights/test_weights(' + 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_weights(' + str(testi) + ').txt'), res_weight, fmt="%.2f")
    end_time = time.time()
    runtime = end_time - start_time
    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'))
    y_train = np.loadtxt(os.path.join(default_path, 'rescale_train_prices/rescale_train_prices(' + str(testi) + ').txt'))
    z_train = np.loadtxt(os.path.join(default_path, 'train_weights/train_weights(' + str(testi) + ').txt'))

    tree_reg_value = tree.DecisionTreeRegressor()
    tree_reg_value.fit(x_train, y_train)
    tree_reg_weight = KNeighborsRegressor(k)
    tree_reg_weight.fit(x_train, z_train)

    x_test = np.loadtxt(os.path.join(default_path, 'test_features/test_features(' + str(testi) + ').txt'))
#    x_test = np.loadtxt('./CombOptNet/test_featureAll.txt')
    y_predict = tree_reg_value.predict(x_test)
    y_real = np.loadtxt(os.path.join(default_path, 'rescale_test_prices/rescale_test_prices(' + str(testi) + ').txt'))
#    y_real = np.loadtxt('./CombOptNet/test_weightsAll.txt')
    res_value = np.vstack((y_real, y_predict))
    res_value = res_value.T
    np.savetxt(os.path.join(default_path, 'CART/CART_prices(' + str(testi) + ').txt'), res_value, fmt="%.2f")
    z_predict = tree_reg_weight.predict(x_test)
    z_real = np.loadtxt(os.path.join(default_path, 'test_weights/test_weights(' + 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_weights(' + str(testi) + ').txt'), res_weight, fmt="%.2f")
    end_time = time.time()
    runtime = end_time - start_time
    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'))
    y_train = np.loadtxt(os.path.join(default_path, 'rescale_train_prices/rescale_train_prices(' + str(testi) + ').txt'))
    z_train = np.loadtxt(os.path.join(default_path, 'train_weights/train_weights(' + str(testi) + ').txt'))

    rf_reg_value = ensemble.RandomForestRegressor()
    rf_reg_value.fit(x_train, y_train)
    rf_reg_weight = KNeighborsRegressor(k)
    rf_reg_weight.fit(x_train, z_train)

    x_test = np.loadtxt(os.path.join(default_path, 'test_features/test_features(' + str(testi) + ').txt'))
#    x_test = np.loadtxt('./CombOptNet/test_featureAll.txt')
    y_predict = rf_reg_value.predict(x_test)
    y_real = np.loadtxt(os.path.join(default_path, 'rescale_test_prices/rescale_test_prices(' + str(testi) + ').txt'))
#    y_real = np.loadtxt('./CombOptNet/test_weightsAll.txt')
    res_value = np.vstack((y_real, y_predict))
    res_value = res_value.T
    np.savetxt(os.path.join(default_path, 'RF/RF_prices(' + str(testi) + ').txt'), res_value, fmt="%.2f")
    z_predict = rf_reg_weight.predict(x_test)
    z_real = np.loadtxt(os.path.join(default_path, 'test_weights/test_weights(' + 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_weights(' + str(testi) + ').txt'), res_weight, fmt="%.2f")
    end_time = time.time()
    runtime = end_time - start_time
    print("training time: ", runtime)
