


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




data_generation_model = 'mismatch'#'iid' 
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!')
#

# Uncomment one of the following if need to experiment with different oracles:
# PATH is the placeholder for the directory to the folder on your computer

# problem_params['oracle_pickle_dir'] = 'PATH/oracle_data/nn/2024-05-22_10-58-32'
# problem_params['regularization_rho'] = 0.1
# problem_params['oracle_pickle_dir'] = 'PATH/oracle_data/nn/2024-05-22_12-04-16'
# problem_params['regularization_rho'] = 0.001

# problem_params['oracle_pickle_dir'] = 'PATH/oracle_data/nn/2024-05-26_12-28-44' # old pwsize = 10
# problem_params['oracle_pickle_dir'] = 'PATH/oracle_data/nn/2024-05-24_11-46-55' # new pwsize = 1000
# problem_params['oracle_pickle_dir'] = 'PATH/oracle_data/nn/2024-05-24_11-50-43' # new pwsize = 100




# print("config", problem_params)
# config_file = open("config_mismatch.pkl", 'rb')
# problem_params = pickle.load(config_file)
#print(problem_params)
# problem_params['test_sample_size'] = 1000
# problem_params['test_batch'] = 1000
# problem_params['n_trails'] = 5

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'] #[0,1,2,3,4]
#deg_grid = [2,3]
if n_input is not None:
    n_grid = [int(n_input)]
else:
    n_grid =  [500] #
#iceo_name = os.path.join(outdir, 'iceo_' + suffix +'.pkl')
#iceo_list = []
for  deg in deg_grid:#[100, 500, 1000]:
    for n_train in n_grid:
        problem_params['train_sample_size'] = n_train
        problem_params['epoch_num'] = 14000
        problem_params['deg_data'] = deg
        # problem_params['regularization_rho'] = 0.001
        problem_params['baseline'] = None
        
        problem_params['entropy_loss'] = True
        problem_params['iceo_reg'] = False
        if data_generation_model == 'iid':
            problem_params['learning_rate'] = 1e-5 #2e-4 + ((n_train-500)/500)*0.7*1e-4 + 0.5*1e-3
        elif data_generation_model == 'mismatch':
            problem_params['learning_rate'] = 1e-3+ ((n_train-500)/500)*0.9*1e-4 #1e-4 + ((n_train-500)/500)*0.9*1e-4 relu
            if n_train < 100:
                problem_params['learning_rate'] = 5e-5
        elif data_generation_model == 'multi_layer':
            problem_params['learning_rate'] = 1e-3+ ((n_train-500)/500)*0.9*1e-4 #1e-4 + ((n_train-500)/500)*0.9*1e-4 relu
        print("learning rate", problem_params['learning_rate'])
        test_results = test_tools.newsvendor_test(problem_params)
        fullname = os.path.join(outdir, 'iceo_' + suffix + 'ntrain' + str(n_train)+ 'deg' + str(deg) + 'entropy'+ str(problem_params['entropy_loss'])+ '.csv')
        file = open(fullname, 'wb')
        test_results.to_csv(fullname)

        problem_params['entropy_loss'] = False
        problem_params['iceo_reg'] = False
        if data_generation_model == 'iid':
            problem_params['learning_rate'] = 2e-4 + ((n_train-500)/500)*0.7*1e-4 + 0.5*1e-3
        elif data_generation_model == 'shift':
            problem_params['learning_rate'] = 5e-5
        elif data_generation_model == 'mismatch':
            problem_params['learning_rate'] = 1e-3+ ((n_train-500)/500)*0.9*1e-4 #1e-4 + ((n_train-500)/500)*0.9*1e-4 relu
            if n_train < 100:
                problem_params['learning_rate'] = 5e-5
        print("learning rate", problem_params['learning_rate'])
        test_results = test_tools.newsvendor_test(problem_params)
        fullname = os.path.join(outdir, 'iceo_' + suffix + 'ntrain' + str(n_train)+ 'deg' + str(deg) + 'entropy'+ str(problem_params['entropy_loss'])+ '.csv')
        file = open(fullname, 'wb')
        test_results.to_csv(fullname)


# for decision reg ICEO
        # problem_params['iceo_reg'] = True
        # problem_params['entropy_loss'] = False
        # if data_generation_model == 'iid':
        #     problem_params['learning_rate'] = 2e-4 + ((n_train-500)/500)*0.7*1e-4 + 0.5*1e-3
        # elif data_generation_model == 'shift':
        #     problem_params['learning_rate'] = 5e-5
        # elif data_generation_model == 'mismatch':
        #     #problem_params['oracle_kernel'] ='polynomial'
        #     #print( problem_params['oracle_kernel'])
        #     problem_params['learning_rate'] = 1e-3+ ((n_train-500)/500)*0.9*1e-4 #1e-4 + ((n_train-500)/500)*0.9*1e-4 relu
        #     if n_train < 100:
        #         problem_params['learning_rate'] = 5e-5
        # print("learning rate", problem_params['learning_rate'])
        # test_results = test_tools.newsvendor_test(problem_params)
        # fullname = os.path.join(outdir, 'iceo_' + suffix + 'ntrain' + str(n_train)+ 'deg' + str(deg) + 'entropy'+ str(problem_params['entropy_loss'])+'reg'+str(problem_params['iceo_reg'])+'.csv')
        # file = open(fullname, 'wb')
        # test_results.to_csv(fullname)






