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 = '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)
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)

x_grid = [5]
deg_grid = [1,2]# problem_params['deg_grid']
if n_input is not None:
    n_grid = [int(n_input)]
else:
    n_grid = problem_params['train_grid']
#iceo_name = os.path.join(outdir, 'iceo_' + suffix +'.pkl')
#iceo_list = []
for  deg in deg_grid:
    for n_train in n_grid:
        for x_dim in x_grid:
            problem_params['train_sample_size'] = n_train
            problem_params['epoch_num'] = 14000
            problem_params['deg_data'] = deg
            problem_params['baseline'] = None
            problem_params['entropy_loss'] = False
            problem_params['learning_rate'] = 1e-4 #8e-4 + ((n_train-500)/500)*1e-4
            problem_params['dim_features']=x_dim
            if data_generation_model == 'mismatch':
                problem_params['learning_rate'] = 1e-4
            test_results = test_tools.network_entropy_test(problem_params)
            #entropy_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')
            #outdir = 'iceo_results/tuning'
            if not os.path.exists(outdir):
                os.mkdir(outdir)
            fullname = os.path.join(outdir, 'entropy_' + suffix + 'ntrain' + str(n_train) + 'deg' + str(deg)+'.csv')
            file = open(fullname, 'wb')
            test_results.to_csv(fullname)
