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' # 'mismatch'#
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 = problem_params['deg_grid']
deg_grid = [1,2,3]
if n_input is not None:
    n_grid = [int(n_input)]
else:
    n_grid = problem_params['train_grid']

for  deg in deg_grid:#[100, 500, 1000]:
    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['dim_features']=x_dim
            if data_generation_model == 'iid':

                problem_params['learning_rate'] = 1e-4 + ((n_train-500)/500)*0.5*1e-4 
            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-4
            print("learning rate", problem_params['learning_rate'])
            test_results = test_tools.network_test(problem_params)
            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, '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'] = True
            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 == 'shift':
                problem_params['learning_rate'] = 5e-5
            elif data_generation_model == 'multi_layer':
                #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
            print("learning rate", problem_params['learning_rate'])
            test_results = test_tools.network_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)
    
