import os
default_n_threads = 8
os.environ['OPENBLAS_NUM_THREADS'] = f"{default_n_threads}"
os.environ['MKL_NUM_THREADS'] = f"{default_n_threads}"
os.environ['OMP_NUM_THREADS'] = f"{default_n_threads}"


from sklearn.ensemble import RandomForestClassifier
import numpy as np
from tools import optimization_oracle_tools

def z_idx(problem_params, z_data):
    size = np.shape(z_data)[0]
    dim_scenario = problem_params['dim_scenario']
    scenario_list = problem_params['scenario_list']
    z_index = np.zeros(size)
    for i in range(size):
        for k in range(dim_scenario):
            # print(z_real[i].numpy())
            # print(scenario_list[k])
            if (z_data[i] == scenario_list[k, :]).all():
                z_index[i] = k
    return z_index
# Fit a classification tree
def get_rf_weights(params, x_train, z_train, x_new, depth = 3, n_trees = 50, min_sample_leaf = 3):
    K = params['dim_scenario']
    z_train_idx = z_idx(params, z_train).astype(int)
    clf = RandomForestClassifier(max_depth = depth, n_estimators = n_trees, min_samples_split = min_sample_leaf)
    clf.fit(x_train, z_train_idx)
    weighted_empirical_distribution = np.zeros(K) 
    for tree in clf.estimators_:
        # Find the leaf node for x_new in the current tree
        leaf_node = tree.apply(x_new)[0]

        # Find all training instances that fall into the same leaf node
        indices = np.where(tree.apply(x_train) == leaf_node)[0]

        # Compute the empirical distribution of y_i for those instances
        empirical_distribution = np.zeros(K)
        for idx in indices:
            empirical_distribution[z_train_idx[idx]] += 1
        empirical_distribution /= len(indices)
        weighted_empirical_distribution += empirical_distribution / len(clf.estimators_)
    return weighted_empirical_distribution



def prescriptive_rf_batch(params, x_train, z_train, x_test, z_test, min_leaf_size = 8, n_trees = 100, depth = 3):
    K = params['dim_scenario']
    dim_sol = params['dim_sol']
    n_test = np.shape(x_test)[0]
    w_rf_batch = np.zeros((n_test, dim_sol))
    for i in range(n_test):
        prob_vec = get_rf_weights(params, x_train, z_train, np.array([x_test[i]]), depth = depth, n_trees = n_trees, min_sample_leaf = min_leaf_size)
        w_rf_batch[i] = optimization_oracle_tools.opt_oracle(prob_vec, params)
    # print("rf_batch", w_rf_batch)
    return w_rf_batch
