text: import random
  import numpy as np
  import matplotlib.pyplot as plt
  from sklearn.tree import DecisionTreeClassifier
  from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
  from sklearn.model_selection import train_test_split
  from sklearn.base import clone
  from sklearn import metrics
  from sklearn import tree
  from tqdm import tqdm
  
  from sklearn.metrics import log_loss, brier_score_loss
  from joblib import Parallel, delayed
  import multiprocessing
  
  from utils import ece, generate_gaussian_parity
  import argparse
  
  
  def generate_gaussian_parity(n, cov_scale=1, angle_params=None, k=1, acorn=None):
  #     means = [[-1.5, -1.5], [1.5, 1.5], [1.5, -1.5], [-1.5, 1.5]]
      means = [[-1, -1], [1, 1], [1, -1], [-1, 1]]
      blob = np.concatenate(
          [
              np.random.multivariate_normal(
                  mean, cov_scale * np.eye(len(mean)), size=int(n / 4)
              )
              for mean in means
          ]
      )
  
      X = np.zeros_like(blob)
      Y = np.concatenate([np.ones((int(n / 4))) * int(i < 2) for i in range(len(means))])
      X[:, 0] = blob[:, 0] * np.cos(angle_params * np.pi / 180) + blob[:, 1] * np.sin(
          angle_params * np.pi / 180
      )
      X[:, 1] = -blob[:, 0] * np.sin(angle_params * np.pi / 180) + blob[:, 1] * np.cos(
          angle_params * np.pi / 180
      )
      return X, Y.astype(int)
  
  
  def get_tree(method='rf', max_depth=1, n_estimators=1, max_leaf_nodes = None ):
      if method == 'gb':
          rf = GradientBoostingClassifier(max_depth=max_depth, n_estimators=n_estimators, random_state=1514, max_leaf_nodes=max_leaf_nodes, learning_rate=1, criterion='mse')
      else:
          rf = RandomForestClassifier(bootstrap=False, max_depth=max_depth, n_estimators=n_estimators, random_state=1514, max_leaf_nodes=max_leaf_nodes)
      
      return rf
  
  def gini_impurity(P1=0, P2=0):
      denom = P1 + P2
      Ginx = 2 * (P1/denom) * (P2/denom)
      return(Ginx)
  def gini_impurity_mean(rf, data, label):
      leaf_idxs = rf.apply(data)
      predict = label
      gini_mean_score = []
      for t in range(leaf_idxs.shape[1]):
          gini_arr = []
          for l in np.unique(leaf_idxs[:,t]):
              cur_l_idx = predict[leaf_idxs[:,t]==l]
              pos_count = np.sum(cur_l_idx)
              neg_count = len(cur_l_idx) - pos_count
              gini = gini_impurity(pos_count, neg_count)
              gini_arr.append(gini) 
  
          gini_mean_score.append(np.array(gini_arr).mean())
      return np.array(gini_mean_score).mean()
  
  def hellinger_explicit(p, q):
      """Hellinger distance between two discrete distributions.
         Same as original version but without list comprehension
      """
      return np.mean(np.sqrt(np.sum((np.sqrt(p) - np.sqrt(q)) ** 2, axis = 1)) / np.sqrt(2))
  
  def pdf(x):
      mu01, mu02, mu11, mu12 = [[-1, -1], [1, 1], [-1, 1], [1, -1]]
     
      cov = 1 * np.eye(2)
      inv_cov = np.linalg.inv(cov) 
  
      p0 = (
          np.exp(-(x - mu01)@inv_cov@(x-mu01).T) 
          + np.exp(-(x - mu02)@inv_cov@(x-mu02).T)
      )/(2*np.pi*np.sqrt(np.linalg.det(cov)))
  
      p1 = (
          np.exp(-(x - mu11)@inv_cov@(x-mu11).T) 
          + np.exp(-(x - mu12)@inv_cov@(x-mu12).T)
      )/(2*np.pi*np.sqrt(np.linalg.det(cov)))
  
      return [p1/(p0+p1), p0/(p0+p1)]
  
  
  def rf_dd_exp(N=4096, reps=100, max_node=None, n_est=10, exp_alias = "depth"):
      
      xx, yy = np.meshgrid(np.arange(-2, 2, 4 / 100), np.arange(-2, 2, 4 / 100))
      true_posterior = np.array([pdf(x) for x in (np.c_[xx.ravel(), yy.ravel()])])
      
      train_mean_error, test_mean_error = [], []
      train_mean_error_log, test_mean_error_log = [], []
      gini_train_mean_score, gini_test_mean_score = [], []
  
      # X, y = get_sample(N)
      X_train, y_train = generate_gaussian_parity(n=N, angle_params=0)
      X_test, y_test = generate_gaussian_parity(n=N, angle_params=0)
  
      method = 'rf'
  
      if max_node is None:
          rf = get_tree(method, max_depth=None)
          rf.fit(X_train, y_train)
          if method == 'gb':
              max_node = (sum([estimator[0].get_n_leaves() for estimator in rf.estimators_])) + 50
          else:
              max_node = sum([estimator.get_n_leaves() for estimator in rf.estimators_]) +50
  
      train_error, test_error  = [list() for _ in range(reps)], [list() for _ in range(reps)]
      train_error_log, test_error_log  = [list() for _ in range(reps)], [list() for _ in range(reps)]
      gini_score_train, gini_score_test = [list() for _ in range(reps)], [list() for _ in range(reps)]
      ece_error = [list() for _ in range(reps)]
      nodes = [list() for _ in range(reps)]
      polys = [list() for _ in range(reps)]
      # for depth in tqdm(range(1, max_node + n_est), position=0, leave=True):
      # for rep_i in tqdm(range(reps), position=0, leave=True):
      def one_run(rep_i): 
          print(rep_i) 
          X_train, y_train = generate_gaussian_parity(n=N, angle_params=0)
          X_test, y_test = generate_gaussian_parity(n=1000, angle_params=0)
  
          rf = get_tree(method, max_depth=1)
          # for rep_i in range(reps):
          for depth in tqdm(range(1, max_node + n_est), position=0, leave=True):
  
              if depth < max_node:
                  rf.max_depth += 1
              else:
                  rf.n_estimators += 3
                  rf.max_depth += 15
                  # rf.warm_start=True
  
              rf.fit(X_train, y_train)
  
              if method == 'gb':
                  nodes[rep_i].append(sum([(estimator[0].get_n_leaves()) for estimator in rf.estimators_])) 
              else:
                  nodes[rep_i].append(sum([estimator.get_n_leaves() for estimator in rf.estimators_]))
              leaf_idxs = rf.apply(X_train)
              polys[rep_i].append(len(np.unique(leaf_idxs)))
              gini_score_train[rep_i].append(gini_impurity_mean(rf, X_train, y_train))
              gini_score_test[rep_i].append(gini_impurity_mean(rf, X_test, y_test))
              train_error[rep_i].append(1 - rf.score(X_train, y_train))
              test_error[rep_i].append(1 - rf.score(X_test, y_test))
              train_error_log[rep_i].append(log_loss(y_train, rf.predict(X_train)))
              test_error_log[rep_i].append(log_loss(y_test, rf.predict(X_test)))
  #             ece_error[rep_i].append( brier_score_loss(y_train, rf.predict(X_train)))
              rf_posteriors_grid = rf.predict_proba(np.c_[xx.ravel(), yy.ravel()])
              ece_error[rep_i].append(hellinger_explicit(rf_posteriors_grid, true_posterior))
          nodes[rep_i] = np.array(nodes[rep_i])
          polys[rep_i] = np.array(polys[rep_i])
          train_error[rep_i] = np.array(train_error[rep_i])
          test_error[rep_i] = np.array(test_error[rep_i])
          train_error_log[rep_i] = np.array(train_error_log[rep_i])
          test_error_log[rep_i] = np.array(test_error_log[rep_i])
          gini_score_train[rep_i] = np.array(gini_score_train[rep_i])
          gini_score_test[rep_i] = np.array(gini_score_test[rep_i])
          ece_error[rep_i] = np.array(ece_error[rep_i])
     
          np.save('results/xor_rf_dd_'+ exp_alias +'_' + str(rep_i)+ ".npy", [nodes[rep_i], polys[rep_i], train_error[rep_i], test_error[rep_i], train_error_log[rep_i], test_error_log[rep_i], gini_score_train[rep_i], gini_score_test[rep_i], ece_error[rep_i]])
          
      num_cores = multiprocessing.cpu_count()
      
      Parallel(n_jobs=-1)(delayed(one_run)(i) for i in range(reps))
  
      train_mean_error = np.array(train_error).mean(axis=0)
      test_mean_error = np.array(test_error).mean(axis=0)
      train_mean_error_log = np.array(train_error_log).mean(axis=0)
      test_mean_error_log = np.array(test_error_log).mean(axis=0)
      nodes_mean = np.array(nodes).mean(axis=0)
      gini_train_mean_score = np.array(gini_score_train).mean(axis=0)
      gini_test_mean_score = np.array(gini_score_test).mean(axis=0)
      error_dict = {'train_err': train_mean_error, 
                  'test_err': test_mean_error,
                  'train_err_log': train_mean_error_log,
                  'test_err_log': test_mean_error_log,
                  'train_gini': gini_train_mean_score,
                  'test_gini': gini_test_mean_score,
                  'nodes': nodes_mean}
      return error_dict
  
  def read_results(reps, exp_alias = "depth"):
      train_error, test_error  = [list() for _ in range(reps)], [list() for _ in range(reps)]
      train_error_log, test_error_log  = [list() for _ in range(reps)], [list() for _ in range(reps)]
      gini_score_train, gini_score_test = [list() for _ in range(reps)], [list() for _ in range(reps)]
      nodes = [list() for _ in range(reps)]
      polys = [list() for _ in range(reps)]
      ece_error = [list() for _ in range(reps)]
      
      for rep_i in range(reps):
          [nodes[rep_i], polys[rep_i], train_error[rep_i], test_error[rep_i], train_error_log[rep_i], test_error_log[rep_i], gini_score_train[rep_i], gini_score_test[rep_i], ece_error[rep_i]] = np.load('results/xor_rf_dd_'+exp_alias+'_' + str(rep_i)+ ".npy")
  
      train_mean_error = np.array(train_error).mean(axis=0)
      test_mean_error = np.array(test_error).mean(axis=0)
      train_mean_error_log = np.array(train_error_log).mean(axis=0)
      test_mean_error_log = np.array(test_error_log).mean(axis=0)
      nodes_mean = np.array(nodes).mean(axis=0)
      polys_mean = np.array(polys).mean(axis=0)
      gini_train_mean_score = np.array(gini_score_train).mean(axis=0)
      gini_test_mean_score = np.array(gini_score_test).mean(axis=0)
      ece_mean_error = np.array(ece_error).mean(axis=0)
      error_dict = {'train_err': train_mean_error, 
                  'test_err': test_mean_error,
                  'train_err_log': train_mean_error_log,
                  'test_err_log': test_mean_error_log,
                  'train_gini': gini_train_mean_score,
                  'test_gini': gini_test_mean_score,
                  ' ece_error': ece_mean_error,
                  'polys': polys_mean,  
                  'nodes': nodes_mean}
  
      return error_dict
  
  # parser = argparse.ArgumentParser(description='Run a double descent experiment.')
  
  # parser.add_argument('--depth', action="store_true", default=False)
  # parser.add_argument('-reps', action="store", dest="reps", type=int)
  # parser.add_argument('-n_est', action="store", dest="n_est", type=int)
  # parser.add_argument('-max_node', action="store", dest="max_node", default=None, type=int)
  # parser.add_argument('-cov_scale', action="store", dest="cov_scale", default=1.0, type=float)
  
  
  # args = parser.parse_args()
  
  # exp_alias = "depth" if args.depth else "width"
  
  # error_dd = rf_dd_exp(max_node=args.max_node, n_est=args.n_est, reps=args.reps, exp_alias = exp_alias)
  