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

warnings.filterwarnings("ignore")
config_file = open("config.pkl", 'rb')
problem_params = pickle.load(config_file)
problem_params['oracle_sample_size'] = 5000
dim_p = problem_params['dim_scenario']
dim_w = problem_params['dim_sol']
outdir = 'oracle_data/pw_dataset'
if not os.path.exists(outdir):
    os.makedirs(outdir)

p_train, p_val, p_test, w_train, w_val, w_test = data_generation_tools_old.pw_dataset(problem_params)
        #print("p_train", p_train)

fullname = os.path.join(outdir, 'p_train' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
pickle.dump(p_train, open(fullname, "wb"))

fullname = os.path.join(outdir, 'p_val' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
pickle.dump(p_val, open(fullname, "wb"))

fullname = os.path.join(outdir, 'p_test' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
pickle.dump(p_test, open(fullname, "wb"))

fullname = os.path.join(outdir, 'w_train' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
pickle.dump(w_train, open(fullname, "wb"))

fullname = os.path.join(outdir, 'w_val' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
pickle.dump(w_val, open(fullname, "wb"))

fullname = os.path.join(outdir, 'w_test' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) + '.pkl')
pickle.dump(w_test, open(fullname, "wb"))
