
import os
import warnings
import pickle
import numpy as np
import torch
import datetime
from tools import data_generation_tools
from tools import optimization_oracle_tools

warnings.filterwarnings("ignore")
config_file = open("/Users/mq56/Library/CloudStorage/Dropbox/E2EL/Numerical/python/python_SGD_multi_NV_20230402/config.pkl", 'rb')
problem_params = pickle.load(config_file)
# load data pickle
dim_p = problem_params['dim_scenario']
dim_w = problem_params['dim_sol']
problem_params['oracle_sample_size'] = 10
outdir = 'oracle_data/pw_dataset/pwsize10'
# outdir = 'oracle_data/pw_dataset'
# outdir = 'oracle_data/pw_dataset/rho3'
# problem_params['regularization_rho'] = 0.001
# rho = problem_params['regularization_rho'] 

fullname = os.path.join(outdir, 'p_train' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
with open(fullname, 'rb') as f:
    p_train = pickle.load(f)

fullname = os.path.join(outdir, 'p_val' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
with open(fullname, 'rb') as f:
    p_val = pickle.load(f)

fullname = os.path.join(outdir, 'p_test' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
with open(fullname, 'rb') as f:
    p_test = pickle.load(f)

fullname = os.path.join(outdir, 'w_train' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
with open(fullname, 'rb') as f:
    w_train = pickle.load(f)

fullname = os.path.join(outdir, 'w_val' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
with open(fullname, 'rb') as f:
    w_val = pickle.load(f)

fullname = os.path.join(outdir, 'w_test' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
with open(fullname, 'rb') as f:
    w_test = pickle.load(f)

print(p_train)

if problem_params['oracle_kernel']=='polynomial':
    oracle_params, mse = optimization_oracle_tools.fit_sol_map_KRR(p_train, w_train, p_test, w_test, problem_params)
    dual_coef, mse = optimization_oracle_tools.dual_coef_KRR(p_train, w_train, p_test, w_test, oracle_params,
                                                                 problem_params)
    print("approx oracle mse", mse)
## dump pickle:
    outdir = 'oracle_data/models'
    dim_p = problem_params['dim_scenario']
    dim_w = problem_params['dim_sol']
    oracle_output = {
        'p_train': p_train,
        'dual_coef': dual_coef,
        'oracle_params': oracle_params,
        'mse': mse
    }

    fullname = os.path.join(outdir, 'oralce_output_' +problem_params['oracle_kernel']+ 'dim_p'+ str(dim_p) + 'dim_w'+ str(dim_w) + '.pkl')
    pickle.dump(oracle_output, open( fullname, "wb" ) )

if problem_params['oracle_kernel'] == 'nn':
    p_train = torch.as_tensor(p_train, dtype=torch.float64)
    w_train = torch.as_tensor(w_train, dtype=torch.float64)
    p_val = torch.as_tensor(p_val, dtype=torch.float64)
    w_val = torch.as_tensor(w_val, dtype=torch.float64)
    p_test = torch.as_tensor(p_test, dtype=torch.float64)
    w_test = torch.as_tensor(w_test, dtype=torch.float64)
    oracle_model, mse = optimization_oracle_tools.fit_sol_map_NN(p_train, w_train, p_val, w_val, p_test, w_test, problem_params, 8e-3, 5000)
    print("approx oracle mape", mse)
    ## dump pickle:
    outdir = 'oracle_data/nn/' + datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
    if not os.path.exists(outdir):
        os.makedirs(outdir)

    #outdir = 'oracle_data/models'
    dim_p = problem_params['dim_scenario']
    dim_w = problem_params['dim_sol']
    oracle_output = {
        'oracle_model': oracle_model,
        'mse': mse
    }

    fullname = os.path.join(outdir, 'oralce_output' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) +'sample'+str(problem_params['oracle_sample_size'])+'.pkl')
    pickle.dump(oracle_output, open(fullname, "wb"))