text: import numpy as np
  import arviz as az
  import seaborn as sns
  import pandas as pd
  import pickle as pkl
  import importlib
  import anndata as ad
  import ast
  
  import matplotlib.pyplot as plt
  from scdcdm.util import result_classes as res
  from scdcdm.util import multi_parameter_sampling as mult
  from scdcdm.util import multi_parameter_analysis_functions as ana
  from scdcdm.util import data_generation as gen
  
  pd.set_option('display.max_columns', 500)
  pd.set_option('display.max_rows', 500)
  sns.set_style("ticks")
  
  #%%
  importlib.reload(ana)
  # Get data
  
  path = "C:\\Users\\Johannes\\Documents\\Uni\\Master's_Thesis\\compositionalDiff-johannes_tests_2\\data\\overall_benchmark"
  
  #%%
  # Find relations between numerical values and composition/increase vectors
  
  b = []
  for y1_0 in [20, 30, 50, 75, 115, 180, 280, 430, 667, 1000]:
      b_i = np.round(gen.counts_from_first(y1_0, 5000, 5), 3)
      b.append(np.round(np.log(b_i / 5000), 2))
  print(b)
  
  b_counts = dict(zip([b_i[0] for b_i in b], [20, 30, 50, 75, 115, 180, 280, 430, 667, 1000]))
  
  b2 = []
  for y1_0 in [20, 30, 50, 75, 115, 180, 280, 430, 667, 1000]:
      b_i = np.round(gen.counts_from_first(y1_0, 5000, 5), 3)
      b2.append(b_i)
  
  b_w_dict = {}
  i = 0
  w_all = []
  for b_i in b2:
      b_t = np.round(np.log(b_i / 5000), 3)
      print(b_t)
      w_d = {}
      for change in [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 200, 400, 600, 800, 1000]:
          _, w = gen.b_w_from_abs_change(b_i, change, 5000)
          w_0 = np.round(w[0], 3)
          w_d[w_0] = change
      b_w_dict[b_t[0]] = w_d
      i += 1
  print(b_w_dict)
  
  #%%
  # all_study_params: One line per data point
  # all_study_params_agg: Aggregate identical simulation parameters
  p = []
  params = []
  for b_i in b:
      print(b_i[0])
      _, all_study_params, all_study_params_agg = ana.multi_run_study_analysis_prepare(path,
                                                                        file_identifier=str(b_i[0]))
  
      all_study_params_agg = ana.get_scores(all_study_params_agg)
      all_study_params = ana.get_scores(all_study_params)
      params.append(all_study_params_agg)
      p.append(all_study_params)
  
  #%%
  # Get some new metrics that make plotting more convenient
  
  all_study_params_agg_2 = pd.concat(params)
  all_study_params_agg_2["n_controls"] = [ast.literal_eval(x)[0] for x in all_study_params_agg_2["n_samples"].tolist()]
  all_study_params_agg_2["n_cases"] = [ast.literal_eval(x)[1] for x in all_study_params_agg_2["n_samples"].tolist()]
  all_study_params_agg_2["n_total"] = all_study_params_agg_2["n_total"].astype("float")
  all_study_params_agg_2["w"] = [ast.literal_eval(x)[0][0] for x in all_study_params_agg_2["w_true"]]
  all_study_params_agg_2["b_0"] = [ast.literal_eval(x)[0] for x in all_study_params_agg_2["b_true"]]
  all_study_params_agg_2["b_count"] = [b_counts[np.round(x, 2)] for x in all_study_params_agg_2["b_0"]]
  
  bs = all_study_params_agg_2["b_0"].tolist()
  ws = all_study_params_agg_2["w"].tolist()
  increases = []
  for i in range(len(bs)):
      increases.append(b_w_dict[bs[i]][ws[i]])
  all_study_params_agg_2["num_increase"] = increases
  all_study_params_agg_2["log_fold_increase"] = np.log2((all_study_params_agg_2["num_increase"] +
                                                         all_study_params_agg_2["b_count"]) /
                                                        all_study_params_agg_2["b_count"])
  
  print(all_study_params_agg_2)
  
  #%%
  # Get some new metrics that make plotting more convenient
  
  all_study_params_2 = pd.concat(p)
  all_study_params_2["n_controls"] = [ast.literal_eval(x)[0] for x in all_study_params_2["n_samples"].tolist()]
  all_study_params_2["n_cases"] = [ast.literal_eval(x)[1] for x in all_study_params_2["n_samples"].tolist()]
  all_study_params_2["n_total"] = all_study_params_2["n_total"].astype("float")
  all_study_params_2["w"] = [ast.literal_eval(x)[0][0] for x in all_study_params_2["w_true"]]
  all_study_params_2["b_0"] = [ast.literal_eval(x)[0] for x in all_study_params_2["b_true"]]
  all_study_params_2["b_count"] = [b_counts[np.round(x, 2)] for x in all_study_params_2["b_0"]]
  
  bs = all_study_params_2["b_0"].tolist()
  ws = all_study_params_2["w"].tolist()
  increases = []
  for i in range(len(bs)):
      increases.append(b_w_dict[bs[i]][ws[i]])
  all_study_params_2["num_increase"] = increases
  all_study_params_2["log_fold_increase"] = np.log2((all_study_params_2["num_increase"] +
                                                         all_study_params_2["b_count"]) /
                                                        all_study_params_2["b_count"])
  
  print(all_study_params_2)
  
  #%%
  # Save that stuff
  result_path = "C:\\Users\\Johannes\\Documents\\Uni\\Master's_Thesis\\compositionalDiff-johannes_tests_2\\data\\benchmark_results"
  
  with open(result_path + "\\normal_results_aggregated.pkl", "wb") as f:
      pkl.dump(all_study_params_agg_2, file=f, protocol=4)
  with open(result_path + "\\normal_results.pkl", "wb") as f:
      pkl.dump(all_study_params_2, file=f, protocol=4)
  
  #%%
  # Exploration...
  sns.heatmap(data=all_study_params_agg_2[["b_count", "num_increase", "mcc"]].pivot_table("mcc", "b_count", "num_increase"), vmin=-1, vmax=1)
  plt.show()
  
  #%%
  
  sns.heatmap(data=all_study_params_agg_2[["n_cases", "n_controls", "mcc"]].pivot_table("mcc", "n_controls", "n_cases"), vmin=-1, vmax=1)
  plt.show()
  
  #%%
  
  print(all_study_params_agg_2[all_study_params_agg_2["n_samples"]=="[10, 10]"])
  
  #%%
  # One heatmap for every combination of b, w
  def draw_heatmap(*args, **kwargs):
      data = kwargs.pop('data')
      d = data.pivot(index=args[1], columns=args[0], values=args[2])
      sns.heatmap(d, **kwargs, vmin=-1, vmax=1)
  
  fg = sns.FacetGrid(all_study_params_agg_2, col='num_increase', row="b_count")
  fg.map_dataframe(draw_heatmap, 'n_controls', 'n_cases', 'mcc', cbar=False)
  plt.show()
  
  #%%
  
  sns.lineplot(data=all_study_params_agg_2, x="b_count", y="mcc")
  plt.show()
  
  #%%
  sns.lineplot(data=all_study_params_agg_2, x="num_increase", y="mcc", hue="b_count")
  plt.show()
  
  #%%
  sns.lineplot(data=all_study_params_agg_2, x="log_fold_increase", y="mcc", hue="b_count")
  plt.show()
  
  #%%
  # Inclusion probability - No results here!
  
  def recalculate_inclusion_probability(results, threshold):
  
      for r in results:
          for res in r.mcmc_results.values():
              res[1]["final_parameter"] = np.where(res[1]["inclusion_prob"] > threshold,
                                                   res[1]["mean_nonzero"],
                                                   0)
          r.get_discovery_rates()
  
      # Generate all_study_params
      all_study_params = pd.concat([r.parameters for r in results])
      simulation_parameters = ["cases", "K", "n_total", "n_samples", "b_true", "w_true", "num_results"]
      all_study_params[simulation_parameters] = all_study_params[simulation_parameters].astype(str)
  
      # Aggregate over identical parameter sets
      all_study_params_agg = all_study_params.groupby(simulation_parameters).sum()
      all_study_params_agg = ana.get_scores(all_study_params_agg.reset_index())
  
      return all_study_params_agg
  
  
  def multiple_threshold_compare(results, thresholds, results_df):
  
      for t in thresholds:
          print("threshold: " + str(t))
          t_results = recalculate_inclusion_probability(results, t)
  
          t_results["threshold"] = t
  
          results_df = results_df.append(t_results)
  
      return results_df
  
  
  import os
  
  
  def multiple_thresholds_with_load(path, thresholds, file_identifier="result_"):
      files = os.listdir(path)
  
      results = []
  
      print("Calculating discovery rates...")
      i = 0
  
      # For all files:
      for f in files:
          i += 1
  
          print("Preparing: ", i / len(files))
          if file_identifier in f:
              # Load file
              r = ana.renamed_load(open(path + "/" + f, "rb"))
  
              results.append(r)
  
      results_df = pd.DataFrame()
  
      results_df = multiple_threshold_compare(results, thresholds, results_df)
  
      return results_df
  
  
  id = "b_[[-2.01 -1.53 -1.53 -1.53 -1.53]]_w_[[[0.034, 0.0, 0.0, 0.0, 0.0]]]"
  results_df = multiple_thresholds_with_load(path, [0.5, 0.7], file_identifier=id)
  
  #%%
  # Calculate metrics for different threshold values
  
  thresholds = np.round(np.arange(0.05, 1, 0.05), 2)
  
  total_df = pd.DataFrame()
  for b_i in b:
      print(b_i[0])
      results_df = multiple_thresholds_with_load(path, thresholds, file_identifier=str(b_i[0]))
  
      total_df = total_df.append(results_df)
  
  #%%
  
  sns.lineplot(data=total_df, x="threshold", y="mcc", hue="")
  plt.show()
  
  #%%
  # extra metrics
  
  total_df_2 = total_df
  total_df_2["n_controls"] = [ast.literal_eval(x)[0] for x in total_df_2["n_samples"].tolist()]
  total_df_2["n_cases"] = [ast.literal_eval(x)[1] for x in total_df_2["n_samples"].tolist()]
  total_df_2["n_total"] = total_df_2["n_total"].astype("float")
  total_df_2["w"] = [ast.literal_eval(x)[0][0] for x in total_df_2["w_true"]]
  total_df_2["b_0"] = [ast.literal_eval(x)[0] for x in total_df_2["b_true"]]
  total_df_2["b_count"] = [b_counts[np.round(x, 2)] for x in total_df_2["b_0"]]
  
  bs = total_df_2["b_0"].tolist()
  ws = total_df_2["w"].tolist()
  increases = []
  for i in range(len(bs)):
      increases.append(b_w_dict[bs[i]][ws[i]])
  total_df_2["num_increase"] = increases
  total_df_2["log_fold_increase"] = np.log2((total_df_2["num_increase"] +
                                             total_df_2["b_count"]) /
                                            total_df_2["b_count"])
  
  print(total_df_2)
  
  #%%
  # Saving
  result_path = "C:\\Users\\Johannes\\Documents\\Uni\\Master's_Thesis\\compositionalDiff-johannes_tests_2\\data\\benchmark_results"
  
  with open(result_path + "\\thresholds.pkl", "wb") as f:
      pkl.dump(total_df_2, file=f, protocol=4)
  
  #%%
  # more plots
  fig, ax = plt.subplots(1, 2, figsize=(13, 5))
  sns.lineplot(data=total_df_2, x="threshold", y="tpr", hue="num_increase", ax=ax[0])
  plt.axvline(x=0.56)
  sns.lineplot(data=total_df_2, x="threshold", y="tnr", hue="num_increase", ax=ax[1])
  plt.axvline(x=0.56)
  
  plt.show()
  
  sns.lineplot(data=total_df_2.loc[total_df_2["num_increase"]>50], x="threshold", y="mcc")
  plt.axvline(x=0.56)
  plt.show()
  
  #%%
  
  # Negative effect testing
  # Same setup as before
  importlib.reload(ana)
  
  path = "C:\\Users\\Johannes\\Documents\\Uni\\Master's_Thesis\\compositionalDiff-johannes_tests_2\\data\\negative_benchmark"
  
  
  _, all_study_params_neg, all_study_params_neg_agg = ana.multi_run_study_analysis_prepare(path)
  
  all_study_params_neg_agg = ana.get_scores(all_study_params_neg_agg)
  all_study_params_neg = ana.get_scores(all_study_params_neg)
  
  
  b = []
  for y1_0 in [115, 280, 1000]:
      b_i = np.round(gen.counts_from_first(y1_0, 5000, 5), 3)
      b.append(np.round(np.log(b_i / 5000), 2))
  print(b)
  
  b_counts = dict(zip([b_i[0] for b_i in b], [115, 280, 1000]))
  
  b2 = []
  for y1_0 in [115, 280, 1000]:
      b_i = np.round(gen.counts_from_first(y1_0, 5000, 5), 3)
      b2.append(b_i)
  
  b_w_dict = {}
  i = 0
  w_all = []
  for b_i in b2:
      b_t = np.round(np.log(b_i / 5000), 3)
      print(b_t)
      w_d = {}
      for change in [-10, -50, -100]:
          _, w = gen.b_w_from_abs_change(b_i, change, 5000)
          w_0 = np.round(w[0], 3)
          w_d[w_0] = change
      b_w_dict[b_t[0]] = w_d
      i += 1
  print(b_w_dict)
  
  all_study_params_neg = all_study_params_neg
  all_study_params_neg["n_controls"] = [ast.literal_eval(x)[0] for x in all_study_params_neg["n_samples"].tolist()]
  all_study_params_neg["n_cases"] = [ast.literal_eval(x)[1] for x in all_study_params_neg["n_samples"].tolist()]
  all_study_params_neg["n_total"] = all_study_params_neg["n_total"].astype("float")
  all_study_params_neg["w"] = [ast.literal_eval(x)[0][0] for x in all_study_params_neg["w_true"]]
  all_study_params_neg["b_0"] = [ast.literal_eval(x)[0] for x in all_study_params_neg["b_true"]]
  all_study_params_neg["b_count"] = [b_counts[np.round(x, 2)] for x in all_study_params_neg["b_0"]]
  
  bs = all_study_params_neg["b_0"].tolist()
  ws = all_study_params_neg["w"].tolist()
  increases = []
  for i in range(len(bs)):
      increases.append(b_w_dict[bs[i]][ws[i]])
  all_study_params_neg["num_increase"] = increases
  all_study_params_neg["log_fold_increase"] = np.log2((all_study_params_neg["num_increase"] +
                                                         all_study_params_neg["b_count"]) /
                                                        all_study_params_neg["b_count"])
  
  print(all_study_params_neg)
  
  all_study_params_neg_agg["n_controls"] = [ast.literal_eval(x)[0] for x in all_study_params_neg_agg["n_samples"].tolist()]
  all_study_params_neg_agg["n_cases"] = [ast.literal_eval(x)[1] for x in all_study_params_neg_agg["n_samples"].tolist()]
  all_study_params_neg_agg["n_total"] = all_study_params_neg_agg["n_total"].astype("float")
  all_study_params_neg_agg["w"] = [ast.literal_eval(x)[0][0] for x in all_study_params_neg_agg["w_true"]]
  all_study_params_neg_agg["b_0"] = [ast.literal_eval(x)[0] for x in all_study_params_neg_agg["b_true"]]
  all_study_params_neg_agg["b_count"] = [b_counts[np.round(x, 2)] for x in all_study_params_neg_agg["b_0"]]
  
  bs = all_study_params_neg_agg["b_0"].tolist()
  ws = all_study_params_neg_agg["w"].tolist()
  increases = []
  for i in range(len(bs)):
      increases.append(b_w_dict[bs[i]][ws[i]])
  all_study_params_neg_agg["num_increase"] = increases
  all_study_params_neg_agg["log_fold_increase"] = np.log2((all_study_params_neg_agg["num_increase"] +
                                                         all_study_params_neg_agg["b_count"]) /
                                                        all_study_params_neg_agg["b_count"])
  
  print(all_study_params_neg_agg)
  
  #%%
  # Save that stuff
  result_path = "C:\\Users\\Johannes\\Documents\\Uni\\Master's_Thesis\\compositionalDiff-johannes_tests_2\\data\\benchmark_results"
  
  with open(result_path + "\\results_negative.pkl", "wb") as f:
      pkl.dump(all_study_params_neg, file=f, protocol=4)
  with open(result_path + "\\results_negative_aggregated.pkl", "wb") as f:
      pkl.dump(all_study_params_neg_agg, file=f, protocol=4)
  #%%
  # Plots
  result_path = "C:\\Users\\Johannes\\Documents\\Uni\\Master's_Thesis\\compositionalDiff-johannes_tests_2\\data\\benchmark_results"
  
  
  def draw_heatmap(*args, **kwargs):
      data = kwargs.pop('data')
      d = data.pivot(index=args[1], columns=args[0], values=args[2])
      sns.heatmap(d, **kwargs, vmin=-1, vmax=1)
  
  fg = sns.FacetGrid(all_study_params_neg_agg, col='num_increase', row="b_count")
  fg.map_dataframe(draw_heatmap, 'n_controls', 'n_cases', 'mcc', cbar=False)
  plt.savefig(result_path + "\\negative_heatmaps.png")
  plt.show()
  
  #%%
  
  # For comparison: Positive heatmaps
  result_path = "C:\\Users\\Johannes\\Documents\\Uni\\Master's_Thesis\\compositionalDiff-johannes_tests_2\\data\\benchmark_results"
  
  with open(result_path + "\\results_aggregated.pkl", "rb") as f:
      all_study_params_agg_pos = pkl.load(file=f)
  
  #%%
  
  fg = sns.FacetGrid(all_study_params_agg_pos.loc[(all_study_params_agg_pos["b_count"].isin([115, 280, 1000])) &
                                                (all_study_params_agg_pos["num_increase"].isin([10, 50, 100]))],
                     col='num_increase', row="b_count")
  fg.map_dataframe(draw_heatmap, 'n_controls', 'n_cases', 'mcc', cbar=False)
  plt.savefig(result_path + "\\pos_heatmaps_for_negative.png")
  plt.show()
  
  # -> For same absolute cell count change, negative effects are detected even a little better!
  
  #%%
  
  def draw_heatmap(*args, **kwargs):
      data = kwargs.pop('data')
      d = data.pivot(index=args[1], columns=args[0], values=args[2])
      sns.heatmap(d, **kwargs, vmin=-1, vmax=1)
  
  fg = sns.FacetGrid(all_study_params_agg_pos.loc[(all_study_params_agg_pos["b_count"].isin([50, 180, 430])) &
                                                (all_study_params_agg_pos["num_increase"].isin([40, 80, 200]))],
                     col='num_increase', row="b_count")
  fg.map_dataframe(draw_heatmap, 'n_controls', 'n_cases', 'mcc', cbar=False)
  plt.savefig(result_path + "\\negative_heatmaps.png")
  plt.show()
  
  