"""
The code is released exclusively for review purposes with the following terms:
PROPRIETARY AND CONFIDENTIAL. UNAUTHORIZED USE, COPYING, OR DISTRIBUTION OF THE 
CODE, VIA ANY MEDIUM, IS STRICTLY PROHIBITED. BY ACCESSING THE CODE, THE 
REVIEWERS AGREE TO DELETE THEM FROM ALL MEDIA AFTER THE REVIEW PERIOD IS OVER.
"""

""" Preprocess and partition the data """
import sys
sys.path.append("../utilities/")
import yaml
import os
import pickle

# Import the necessary modules
from data_processing import meps_data, iris_data
from utils import fname_data, create_dir_if_not_exist

# Pass arguments and run the code
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--config_fname")
parser.add_argument("--dataset_key")
args = parser.parse_args()

data_preproc_functions = {"MEPS": meps_data,
                          "IRIS": iris_data}

# Load the config file
config = yaml.load(open(
            os.path.join("config", args.config_fname)),
            Loader=yaml.FullLoader)

# create the input dataset
dataset = data_preproc_functions[args.dataset_key](
            **config["Data_Preproc"][args.dataset_key])
# dataset is a tuple that contains:
# ((X_train, X_test, y_train, y_test, w_train, w_test),
#             categorical_feature_names, numerical_feature_names,
#             categorical_feature_inds, numerical_feature_inds,
#             colnames_onehot, colnames_orig) 

datafname = fname_data(config, args.dataset_key)+".pkl"
dirname = os.path.join("data", args.dataset_key, "input")
create_dir_if_not_exist(dirname)
pickle.dump(dataset, open(os.path.join(dirname, datafname), "wb"))
