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")



data_generation_model = 'mismatch'#''iid'#'mismatch'#'sys.argv[1] 
n_input = None # sys.argv[2]
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)
else:
    raise Exception(
        'Invalid data generation model!')
#
print("config", problem_params)

suffix = problem_params['data_generation_model']
outdir = 'iceo_results/'+ datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
if not os.path.exists(outdir):
    os.mkdir(outdir)


deg_grid = problem_params['deg_grid'] 
n_grid = problem_params['train_grid']
if n_input is not None:
    n_grid = [int(n_input)]
else:
    n_grid = problem_params['train_grid']
#deg_grid = [2]
for  deg in deg_grid:
    for n_train in n_grid:#[100, 500, 1000]:
        problem_params['train_sample_size'] = n_train
        problem_params['deg_data'] = deg
        problem_params['epoch_num'] = 14000
        problem_params['learning_rate'] = 1e-3 + ((n_train-500)/500)*1e-4
        if data_generation_model == 'mismatch':
            problem_params['learning_rate'] = 1e-4
        test_results = test_tools.newsvendor_entropy_test(problem_params)
        fullname = os.path.join(outdir, 'entropy_' + suffix + 'ntrain' + str(n_train) + 'deg' + str(deg)+'.csv')
        file = open(fullname, 'wb')
        test_results.to_csv(fullname)


