//  [python]
def plot_bucket_ite_comparison(treatment_effect_tracker):
  """Compare ITE values for diff-performance models."""
  # unique_hparam_types = treatment_effect_tracker['hparam_type'].unique()
  # for hparam_type in unique_hparam_types:
  for hparam_type in ['config.optimizer']:
    not_range_accuracy = '0.0_1.0'
    kernel_type = 'rbf'
    filtered = treatment_effect_tracker.where(
        (treatment_effect_tracker['hparam_type'] == hparam_type) &
        (treatment_effect_tracker['kernel_type'] == kernel_type) &
        (treatment_effect_tracker['range_accuracy'] != not_range_accuracy)
    ).dropna()
    for target_type in ['y_pred', 'explan']:
      g = sns.catplot(
          x='range_accuracy',
          y='%s_ite' % target_type,
          hue='h1_h2_str',
          col=None if target_type == 'y_pred' else 'explan_type',
          data=filtered,
          kind='violin',
          sharey='row',
          legend=False,
          # height=3,
          linewidth=0.6,
      )
      if target_type == 'y_pred':
        g.fig.set_size_inches(4, 4)
      else:
        g.fig.set_size_inches(16, 4)
      for ax_idx, ax in enumerate(g.axes[0]):
        ax.set_xticklabels(ax.get_xticklabels(), rotation=30)
        if ax_idx == 0 and target_type == 'explan':
          ax.legend(
              loc='upper left',
              ncol=1,
              title=hparam_type,
          )

      for ax_idx, ax in enumerate(g.axes.flatten()):
        ax.set_xlabel('Test Accuracy', fontsize=14)
        if target_type == 'y_pred':
          if ax_idx == 0: ax.set_ylabel(r'$ITE_{Y}(x)$', fontsize=14)
        else:
          explan_type = ax.get_title().split('=')[1][1:]
          explan_type = config.EXPLAN_NAME_CONVERTER[explan_type]
          ax.set_title(explan_type)
          if ax_idx == 0: ax.set_ylabel(r'$ITE_{E}(x)$', fontsize=14)

      for ax_idx, ax in enumerate(g.axes.flatten()):
        labels = ax.get_xticklabels()
        for label in labels:
          range_accuracy = label.get_text()
          range_percentile = config.RANGE_ACCURACY_CONVERTER[range_accuracy]
          label.set_text(range_percentile)
        ax.set_xticklabels(labels)

      plt.subplots_adjust(hspace=0.1, wspace=0.1)
      save_paper_figures(
          'bucket_ite_comparison_%s_%s_%s_%s.png' %
          (config.cfg.DATASET, target_type, kernel_type, hparam_type)
      )

  ##############################################################################
  # repeat for all hparams in 1 fig
  ##############################################################################

  not_range_accuracy = '0.0_1.0'
  kernel_type = 'rbf'
  filtered = treatment_effect_tracker.where(
      (treatment_effect_tracker['kernel_type'] == kernel_type) &
      (treatment_effect_tracker['range_accuracy'] != not_range_accuracy)
  ).dropna()

  unique_hparam_types = treatment_effect_tracker['hparam_type'].unique()
  unique_explan_types = treatment_effect_tracker['explan_type'].unique()

  num_rows = len(unique_hparam_types)
  num_cols = len(unique_explan_types) + 1

  fig, axes = plt.subplots(num_rows, num_cols, sharex=True, sharey='row')

  for row_idx in range(num_rows):
    for col_idx in range(num_cols):
      ax = axes[row_idx, col_idx]
      hparam_type = unique_hparam_types[row_idx]

      if col_idx == 0:
        y = 'y_pred_ite'
        data = filtered.where(
            (filtered['hparam_type'] == hparam_type)
        ).dropna()
      else:
        y = 'explan_ite'
        explan_type = unique_explan_types[col_idx - 1]
        data = filtered.where(
            (filtered['hparam_type'] == hparam_type) &
            (treatment_effect_tracker['explan_type'] == explan_type)
        ).dropna()

      sns.violinplot(
          x='range_accuracy',
          y=y,
          hue='h1_h2_str',
          data=data,
          ax=ax,
          linewidth=0.4,
          legend=False
      )

      if col_idx == 0:
        ax.set_ylabel(r'$ITE_{Y}(x)$', fontsize=14)
      elif col_idx == 1:
        ax.set_ylabel(r'$ITE_{E}(x)$', fontsize=14)
      else:
        ax.set_ylabel('')

      if row_idx == len(unique_hparam_types) - 1:
        ax.set_xlabel('Test Accuracy', fontsize=14)
        labels = ax.get_xticklabels()
        for label in labels:
          range_accuracy = label.get_text()
          range_percentile = config.RANGE_ACCURACY_CONVERTER[range_accuracy]
          label.set_text(range_percentile)
        ax.set_xticklabels(labels, rotation=30)
      else:
        ax.set_xlabel('')
        ax.set_xticklabels('')

      if col_idx == 2:
        ax.legend(
            loc='upper center',
            bbox_to_anchor=(0.5, 1.18),
            ncol=len(data['h1_h2_str'].unique()),
            title='',
        )
      else:
        ax.legend().set_visible(False)

  plt.plot([0.27, 0.27], [0.10, 0.88], color='lightgray', lw=2,
      transform=plt.gcf().transFigure, clip_on=False)
  fig.set_size_inches(24, 32)
  save_paper_figures(
      'bucket_ite_comparison_%s_%s_all.png' %
      (config.cfg.DATASET, kernel_type)
  )