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)
problem_params['oracle_sample_size'] = 10
dim_p = problem_params['dim_scenario']
dim_w = problem_params['dim_sol']
outdir = 'oracle_data/pw_dataset'
# outdir = 'oracle_data/pw_dataset/rho3'
outdir = 'oracle_data/pw_dataset/pwsize10'
if not os.path.exists(outdir):
    os.makedirs(outdir)

# problem_params['regularization_rho'] = 0.001
p_train, p_val, p_test, w_train, w_val, w_test = data_generation_tools.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"))
