"""
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.
"""

""" Create bootstrap samples from the base perturbations based on the configurations
provided. This will create the different environments. """
import numpy as np
import sys
sys.path.append("../utilities/")
import os

from joblib import Parallel, delayed
from sklearn.utils import check_random_state
import yaml
import pickle

from utils import (fname_base_perts, fname_env_perts, compute_weights, 
                    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")
parser.add_argument("--pert_key")
args = parser.parse_args()

# Load the config file
config = yaml.load(open(
            os.path.join("config", args.config_fname)),
            Loader=yaml.FullLoader)
 
# dump the base perts
base_pert_fname = fname_base_perts(config, 
                    args.pert_key, args.dataset_key)+".pkl"
dirname = os.path.join("data", args.dataset_key, "perturbations")
base_perturbations = pickle.load(open( 
                        os.path.join(dirname, base_pert_fname), "rb" ) )

# create the indices for the environments
n_data_all = len(base_perturbations["indices"])
# Maximum number of perturbations possible
max_n_perts_env = base_perturbations["samp_perts_exp"][0].shape[0]
n_perts_env = config["Env_Perturbations"]["cnt"]
num_envs = config["Env_Perturbations"]["num_envs"]
normalize_weights = config["Env_Perturbations"]["normalize_weights"]
kernel_width = config["Env_Perturbations"]["kernel_width"]
random_state = check_random_state(config["Env_Perturbations"]["random_seed"])
print(max_n_perts_env)

# Sample indices for the environments (bootstraps)
samp_inds_env = np.zeros((n_data_all, n_perts_env, num_envs), dtype=int)
for env in range(num_envs):
    for idx in range(n_data_all):
        samp_inds_env[idx, 1:n_perts_env, env] = random_state.choice(
                                            np.arange(1, max_n_perts_env),
                                            size=n_perts_env-1)

env_perturbations = {"samp_inds_env": samp_inds_env}
                    #  "weights_env": weights_env}

# dump the base perts
env_pert_fname = fname_env_perts(config, "Env_Perturbations",
                    args.pert_key, args.dataset_key)+".pkl"
dirname = os.path.join("data", args.dataset_key, "perturbations")
create_dir_if_not_exist(dirname)
pickle.dump(env_perturbations, open( os.path.join(dirname, env_pert_fname), "wb" ) )

