code: def balanced_train_test_split(
      x: pd.DataFrame,
      y: Union[np.ndarray, list],
      *cross_cols: str,
      categorical_cols: List[str] = [],
      min_per_group: int = 20,
      max_per_group: Optional[int] = None,
      fallback_type: str = "upsample",
      bins: int = 6,
      random_state: int = None,
  ) -> Tuple[
      pd.DataFrame,
      np.ndarray,
      pd.DataFrame,
      np.ndarray,
      np.ndarray,
      np.ndarray,
  ]:
      """
      Splits the "x" DataFrame and "y" Array into train/test split training sets with 
      a balanced number of examples for each of the categories of the columns provided.
      For example, if the columns provided are "gender" and "loan", the resulting splits
      would contain an equal number of examples for Male with Loan Approved, Male with 
      Loan Rejected, Female with Loan Approved, and Female with Loan Rejected. The 
      "fallback_type" parameter provides the behaviour that is triggered if there are not
      enough datapoint examples for one of the subcategory groups - the default is "half"
  
      Example
      -------
  
      .. code-block:: python
  
          x: pd.DataFrame # Contains the input features
          y: np.array # Contains the labels for the data
          categorical_cols: List[str] # Name of columns that are categorical
  
          x_train, y_train, x_test, y_test, train_idx, test_idx = \\
              xai.balanced_train_test_split(
                      x, y, balance_on=["gender"], 
                      categorical_cols=categorical_cols, min_per_group=300,
                      fallback_type="half")
  
      Args
      -----
      x : 
          Pandas dataframe containing all the features in dataset
      y : 
          Array containing "actual" labels for the dataset 
      *cross_cols :
          One or more positional arguments (passed as *args) that 
          are used to split the data into the cross product of their values 
      categorical_cols : 
          [Default: []] Columns within dataframe that are
          categorical. Columns that are not np.objects and are not part explicitly
          provided here will be treated as numeric, and bins will be used.
      min_per_group : 
          [Default: 20] This is the number of examples for each
          of the groups created
      max_per_group : 
          [Default: None] This is the maximum number of examples for
          each group to be provided with.
      fallback_type : 
          [Default: upsample] This is the fallback mechanism for when
          one of the groups contains less elements than the number provided 
          through min_per_group. The options are "upsample", "ignore" and "error". 
              - "upsample": This will get samples with replacement so will repeat elements
              - "ignore": Will just ignore and return all the elements available
              - "error": Throw an exception for any groups with less elements
      bins : 
          [Default: 6] Number of bins to be used for numerical cols
      random_state: 
          [Default: None] Random seed for the internal sampling
  
      Returns
      -------
  
      x_train : pd.DataFrame 
          DataFrame containing traning datapoints
      y_train : np.ndarray 
          Array containing labels for training datapoints
      x_test : pd.DataFrame
          DataFrame containing test datapoints
      y_test : np.ndarray
          Array containing labels for test datapoints
      train_idx : np.ndarray 
          Boolean array with True on Training indexes
      test_idx : np.ndarray 
          Boolean array with True on Testing indexes
  
      """
  
      if not cross_cols:
          raise TypeError(
              "imbalance_plot requires at least 1 string column name"
          )
      if min_per_group < 1:
          raise TypeError("min_per_group must be at least 1")
      if max_per_group and max_per_group < min_per_group:
          raise TypeError(
              f"min_per_group ({min_per_group}) must be less or equal than "
              f"max_per_group ({max_per_group}) if max_per_group is provided."
          )
  
      if random_state:
          random.setstate(random_state)
  
      tmp_df = x.copy()
      tmp_df["target"] = y
  
      cross = ["target"] + list(cross_cols)
  
      if not categorical_cols:
          categorical_cols = _infer_categorical(tmp_df)
  
      # TODO: Enable for non-categorical targets
      categorical_cols = ["target"] + categorical_cols
  
      grouped = group_by_columns(
          tmp_df, cross, bins=bins, categorical_cols=categorical_cols
      )
  
      def resample(x):
          group_size = x.shape[0]
  
          if max_per_group:
              if group_size > max_per_group:
                  return x.sample(max_per_group)
  
          if group_size > min_per_group:
              return x.sample(min_per_group)
  
          if fallback_type == "upsample":
              return x.sample(min_per_group, replace=True)
          elif fallback_type == "ignore":
              return x
          elif fallback_type == "error":
              raise ValueError(
                  "Number of samples for group are not enough,"
                  " and fallback_type provided was 'error'"
              )
          else:
              raise (
                  f"Sampling type provided not found: given {fallback_type}, "
                  "expected: 'error', or 'half'"
              )
  
      group = grouped.apply(resample)
  
      selected_idx = [g[-1] for g in group.index.values]
  
      train_idx = np.full(tmp_df.shape[0], True, dtype=bool)
      train_idx[selected_idx] = False
      test_idx = np.full(tmp_df.shape[0], False, dtype=bool)
      test_idx[selected_idx] = True
  
      df_train = tmp_df.iloc[train_idx]
      df_test = tmp_df.iloc[test_idx]
  
      x_train = df_train.drop("target", axis=1)
      y_train = df_train["target"].values
      x_test = df_test.drop("target", axis=1)
      y_test = df_test["target"].values
  
      return x_train, y_train, x_test, y_test, train_idx, test_idx
  
imports: import pandas
  import matplotlib.pyplot
  import numpy
  from scipy.stats import spearmanr
  from scipy.cluster import hierarchy
  from typing import List, Any, Union, Tuple, Optional, Dict
  import random, math
  from sklearn.metrics import roc_curve, precision_recall_curve, roc_auc_score, confusion_matrix
  import itertools
  import logging
label_desc: Use image node_id 10 for calling a global function with example usage: balanced_train_test_split(x, y) and returns: x_train,  y_train,  x_test,  y_test,  train_idx,  test_idx, x, x, x, x