import os
default_n_threads = 4
os.environ['OPENBLAS_NUM_THREADS'] = f"{default_n_threads}"
os.environ['MKL_NUM_THREADS'] = f"{default_n_threads}"
os.environ['OMP_NUM_THREADS'] = f"{default_n_threads}"


import pandas as pd
import random
import numpy as np
from tools import test_tools
import os
import datetime
import pickle
import sys
import warnings
warnings.filterwarnings("ignore")
#random.seed(465)

data_generation_model = 'iid' # sys.argv[1]
if data_generation_model == 'iid':
    config_file = open("config.pkl", 'rb')
    problem_params = pickle.load(config_file)
elif data_generation_model == 'mismatch':
    config_file = open("config_mismatch.pkl", 'rb')
    problem_params = pickle.load(config_file)
elif data_generation_model == 'shift':
    config_file = open("config_shift.pkl", 'rb')
    problem_params = pickle.load(config_file)
else:
    raise Exception(
        'Invalid data generation model!')
#
print("config", problem_params)

deg_grid = problem_params['deg_grid'] #[0,1,2,3,4]
n_grid = problem_params['train_grid']
problem_params['rf_min_leaf_grid'] = [5]
# n_grid =
# problem_params['baseline'] = [ 'saa', 'knn', 'kernel']
problem_params['baseline'] = ['saa', 'forest']
# baseline_list = []
for  deg in deg_grid:
    # n_train = 500
    for n_train in n_grid:#[100, 500, 1000]:
        for x_dim in [5]:
            problem_params['train_sample_size'] = n_train
            problem_params['deg_data'] = deg
            problem_params['dim_features']=x_dim
            test_results = test_tools.network_test_baseline_only(problem_params)
        #baseline_list.append(test_results)
            suffix = problem_params['data_generation_model']+'x_dim' + str(problem_params['dim_features'])
            outdir = 'iceo_results/'+ datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
            if not os.path.exists(outdir):
                os.mkdir(outdir)
            fullname = os.path.join(outdir, 'forest_' + suffix + 'ntrain'+ str(n_train)+'deg' + str(deg)  +'.csv')
            file = open(fullname, 'wb')
            test_results.to_csv(fullname)
#pickle.dump(baseline_list, file)
