text: from copy import deepcopy
  from pprint import pprint
  
  import matplotlib
  import numpy as np
  import pandas as pd
  from matplotlib import pyplot as plt
  
  from constants import (RESULTS_CF_DISSABLEMENT_FILE,
                         RESULTS_CF_SUFFICIENCY_FILE, RESULTS_OBS_FILE,
                         VIGNETTES_FILE)
  from helpers import bintest, doctor_top_ns, mean_list
  from utils import load_from_json, load_from_pickle
  
  matplotlib.use("Agg")
  
  
  def produce_results(*, args):
      print(f"> Producing results from {str(args.results.absolute())}")
  
      topn_results_obs = load_from_pickle(args.results / RESULTS_OBS_FILE)
      topn_results_counter_diss = load_from_pickle(
          args.results / RESULTS_CF_DISSABLEMENT_FILE
      )
      topn_results_counter_suff = load_from_pickle(
          args.results / RESULTS_CF_SUFFICIENCY_FILE
      )
  
      make_supplementary_table_one(
          args=args,
          topn_results_obs=topn_results_obs,
          topn_results_counter_diss=topn_results_counter_diss,
          topn_results_counter_suff=topn_results_counter_suff,
      )
      make_figure_three(
          args=args,
          topn_results_obs=topn_results_obs,
          topn_results_counter_diss=topn_results_counter_diss,
          topn_results_counter_suff=topn_results_counter_suff,
      )
      make_table_one_and_supplementary_table_two(
          args=args,
          topn_results_obs=topn_results_obs,
          topn_results_counter_diss=topn_results_counter_diss,
          topn_results_counter_suff=topn_results_counter_suff,
      )
      df_results, doc_topn = make_supplementary_table_three(
          args=args,
          topn_results_obs=topn_results_obs,
          topn_results_counter_diss=topn_results_counter_diss,
          topn_results_counter_suff=topn_results_counter_suff,
      )
      make_table_two(args=args, df_results=df_results, doc_topn=doc_topn)
      make_figure_four(args=args, df_results=df_results)
  
  
  def make_supplementary_table_one(
      *, args, topn_results_obs, topn_results_counter_diss, topn_results_counter_suff
  ):
  
      p_obs = sum(topn_results_obs) / len(topn_results_obs)
      error_obs = np.sqrt(p_obs * (1 - p_obs) / len(topn_results_obs))
  
      p_counter = sum(topn_results_counter_suff) / len(topn_results_counter_suff)
      error_counter = np.sqrt(
          p_counter * (1 - p_counter) / len(topn_results_counter_suff)
      )
  
      _p_counter = sum(topn_results_counter_diss) / len(topn_results_counter_diss)
      _error_counter = np.sqrt(
          _p_counter * (1 - _p_counter) / len(topn_results_counter_diss)
      )
  
      dg2 = pd.DataFrame(
          {
              "N": list(np.arange(20) + 1),
              "Posterior Accuracy": list(p_obs),
              "Posterior Error": list(error_obs),
              "Expected sufficiency": list(p_counter),
              "ES error": list(error_counter),
              "Expected disablement": list(_p_counter),
              "ED error": list(_error_counter),
          }
      )
  
      dg2["Posterior"] = (
          dg2["Posterior Accuracy"].round(3).apply(str)
          + " $ \pm $ "
          + dg2["Posterior Error"].round(3).apply(str)
      )
      dg2["Disablement"] = (
          dg2["Expected disablement"].round(3).apply(str)
          + " $ \pm $ "
          + dg2["ED error"].round(3).apply(str)
      )
      dg2["Sufficiency"] = (
          dg2["Expected sufficiency"].round(3).apply(str)
          + " $ \pm $ "
          + dg2["ES error"].round(3).apply(str)
      )
  
      print(
          dg2[["N", "Posterior", "Disablement", "Sufficiency"]].to_latex(
              escape=False, index=False
          )
      )
  
  
  def make_figure_three(
      *, args, topn_results_obs, topn_results_counter_diss, topn_results_counter_suff
  ):
      topn_results_obs = load_from_pickle(args.results / RESULTS_OBS_FILE)
      topn_results_counter_suff = load_from_pickle(
          args.results / RESULTS_CF_SUFFICIENCY_FILE
      )
  
      x = sum(topn_results_obs) / len(topn_results_obs)
      y = sum(topn_results_counter_suff) / len(topn_results_counter_suff)
      plt.figure(figsize=(5, 4))
  
      p_obs = sum(topn_results_obs) / len(topn_results_obs)
      error_obs = np.sqrt(p_obs * (1 - p_obs) / len(topn_results_obs))
      p_counter = sum(topn_results_counter_suff) / len(topn_results_counter_suff)
      error_counter = np.sqrt(
          p_counter * (1 - p_counter) / len(topn_results_counter_suff)
      )
  
      xmarks = [i + 1 for i in range(len(p_obs))]
  
      plt.plot(xmarks, 1 - p_obs, label="Associative", color="blue")
      plt.plot(xmarks, 1 - p_counter, label="Counterfactual", color="seagreen")
      plt.plot(xmarks, 1 - (1 - y) / (1 - x), linestyle="--", color="black")
  
      plt.fill_between(
          xmarks,
          1 - p_obs - 2 * error_obs,
          1 - p_obs + 2 * error_obs,
          alpha=0.2,
          edgecolor="#1B2ACC",
          facecolor="#089FFF",
          linewidth=0,
          linestyle="None",
          antialiased=True,
      )
  
      plt.fill_between(
          xmarks,
          1 - p_counter - 2 * error_counter,
          1 - p_counter + 2 * error_counter,
          alpha=0.2,
          edgecolor="#1B2ACC",
          facecolor="seagreen",
          linewidth=0,
          linestyle="None",
          antialiased=True,
      )
  
      plt.xticks([i + 1 for i in range(16)])
      plt.xlim(1, 15)
      plt.ylim(0, 0.5)
      plt.show()
      plt.savefig(args.results / "algo_vs_algo.pdf")
  
  
  def make_table_one_and_supplementary_table_two(
      *, args, topn_results_obs, topn_results_counter_diss, topn_results_counter_suff
  ):
      casecards = load_from_json(args.datapath / VIGNETTES_FILE)
  
      results_obs = {
          "common": [],
          "rare": [],
          "very_rare": [],
          "almost_impossible": [],
          "uncommon": [],
          "very_common": [],
      }
      results_counter = {
          "common": [],
          "rare": [],
          "very_rare": [],
          "almost_impossible": [],
          "uncommon": [],
          "very_common": [],
      }
      wins_obs = {
          "common": 0,
          "rare": 0,
          "very_rare": 0,
          "almost_impossible": 0,
          "uncommon": 0,
          "very_common": 0,
      }
      wins_counter = {
          "common": 0,
          "rare": 0,
          "very_rare": 0,
          "almost_impossible": 0,
          "uncommon": 0,
          "very_common": 0,
      }
      draws = {
          "common": 0,
          "rare": 0,
          "very_rare": 0,
          "almost_impossible": 0,
          "uncommon": 0,
          "very_common": 0,
      }
  
      for num, card in enumerate(casecards.values()):
          if args.first is not None and num >= args.first:
              continue
  
          rareness = card["card"]["diseases"][0]["rareness"]
          r_obs = sum(topn_results_obs[num])
          r_suff = sum(topn_results_counter_suff[num])
          results_obs[rareness] += [min(21 - r_obs, 20)]
          results_counter[rareness] += [min(21 - r_suff, 20)]
  
          if r_obs > r_suff:
              wins_obs[rareness] += 1
          elif r_obs < r_suff:
              wins_counter[rareness] += 1
          else:
              draws[rareness] += 1
  
      results_obs_sum = []
      wins_obs_all, wins_counter_all, draws_all = 0, 0, 0
      for k, val in results_obs.items():
          results_obs_sum += val
          wins_obs_all += wins_obs[k]
          draws_all += draws[k]
          results_obs[k] = {"mean": np.mean(val), "std": np.std(val)}
      results_obs["all"] = {
          "mean": np.mean(results_obs_sum),
          "std": np.std(results_obs_sum),
      }
  
      results_counter_sum = []
      wins_counter_all = 0
      for k, val in results_counter.items():
          results_counter_sum += val
          wins_counter_all += wins_counter[k]
          results_counter[k] = {"mean": np.mean(val), "std": np.std(val)}
      results_counter["all"] = {
          "mean": np.mean(results_counter_sum),
          "std": np.std(results_counter_sum),
      }
      draws["all"] = draws_all
      wins_obs["all"] = wins_obs_all
      wins_counter["all"] = wins_counter_all
  
      print("> Observational Results")
      pprint(results_obs)
      print("")
  
      print("> Counterfactual Results")
      pprint(results_counter)
      print("")
  
  
  def make_supplementary_table_three(
      *, args, topn_results_obs, topn_results_counter_diss, topn_results_counter_suff
  ):
      casecards = load_from_json(args.datapath / VIGNETTES_FILE)
  
      paired_results = {}
      doc_topn = {}
      doc_topn_caseav = {}
      doc_score, obs_score = [], []
      for num, card in enumerate(casecards.values()):
          if args.first is not None and num >= args.first:
              continue
  
          true_id = card["card"]["diseases"][0]["id"]
          pred_suff = topn_results_counter_suff[num]
          pred_diss = topn_results_counter_diss[num]
          pred_obs = topn_results_obs[num]
          doc_res_n = doctor_top_ns(card, true_id)
          for val in doc_res_n:
              if val[1] == 0:
                  continue
              if val[0] not in paired_results.keys():
                  paired_results[val[0]] = [
                      [
                          deepcopy(val[2]),
                          deepcopy(pred_suff[val[1] - 1]),
                          deepcopy(pred_diss[val[1] - 1]),
                          deepcopy(pred_obs[val[1] - 1]),
                      ]
                  ]
              else:
                  paired_results[val[0]] += [
                      [
                          deepcopy(val[2]),
                          deepcopy(pred_suff[val[1] - 1]),
                          deepcopy(pred_diss[val[1] - 1]),
                          deepcopy(pred_obs[val[1] - 1]),
                      ]
                  ]
          for val in doc_res_n:
              if val[1] == 0:
                  continue
  
              if val[0] not in doc_topn.keys():
                  doc_topn[val[0]] = {
                      "count": 1,
                      "sufficiency": {
                          val[1]: np.array([1, deepcopy(pred_suff[val[1] - 1])])
                      },
                      "disablement": {
                          val[1]: np.array([1, deepcopy(pred_diss[val[1] - 1])])
                      },
                      "obs": {val[1]: np.array([1, deepcopy(pred_obs[val[1] - 1])])},
                      "doctor": {val[1]: np.array([1, deepcopy(val[2])])},
                  }
              else:
                  doc_topn[val[0]]["count"] += 1
                  if (
                      val[1] not in doc_topn[val[0]]["sufficiency"].keys()
                  ):  # this doctor has never had this score before
  
                      doc_topn[val[0]]["sufficiency"][val[1]] = np.array(
                          [1, deepcopy(pred_suff[val[1] - 1])]
                      )
                      doc_topn[val[0]]["disablement"][val[1]] = np.array(
                          [1, deepcopy(pred_diss[val[1] - 1])]
                      )
                      doc_topn[val[0]]["obs"][val[1]] = np.array(
                          [1, deepcopy(pred_obs[val[1] - 1])]
                      )
                      doc_topn[val[0]]["doctor"][val[1]] = np.array([1, deepcopy(val[2])])
                  else:
                      doc_topn[val[0]]["sufficiency"][val[1]] += np.array(
                          [1, deepcopy(pred_suff[val[1] - 1])]
                      )
                      doc_topn[val[0]]["disablement"][val[1]] += np.array(
                          [1, deepcopy(pred_diss[val[1] - 1])]
                      )
                      doc_topn[val[0]]["obs"][val[1]] += np.array(
                          [1, deepcopy(pred_obs[val[1] - 1])]
                      )
                      doc_topn[val[0]]["doctor"][val[1]] += np.array(
                          [1, deepcopy(val[2])]
                      )
  
          this_card_res_doc = {
              1: [],
              2: [],
              3: [],
              4: [],
              5: [],
              6: [],
              7: [],
              8: [],
              9: [],
          }
          this_card_res_suff = {
              1: [],
              2: [],
              3: [],
              4: [],
              5: [],
              6: [],
              7: [],
              8: [],
              9: [],
          }
          this_card_res_diss = {
              1: [],
              2: [],
              3: [],
              4: [],
              5: [],
              6: [],
              7: [],
              8: [],
              9: [],
          }
          this_card_res_obs = {
              1: [],
              2: [],
              3: [],
              4: [],
              5: [],
              6: [],
              7: [],
              8: [],
              9: [],
          }
          for val in doc_res_n:
              if val[1] == 0:
                  continue
              if val[1] > 9:
                  continue
              this_card_res_doc[val[1]] += [val[2]]
              this_card_res_suff[val[1]] += [deepcopy(pred_suff[val[1] - 1])]
              this_card_res_diss[val[1]] += [deepcopy(pred_diss[val[1] - 1])]
              this_card_res_obs[val[1]] += [deepcopy(pred_obs[val[1] - 1])]
          this_card_res_doc = dict(
              [[k, mean_list(val)] for k, val in this_card_res_doc.items()]
          )
          this_card_res_suff = dict(
              [[k, mean_list(val)] for k, val in this_card_res_suff.items()]
          )
          this_card_res_diss = dict(
              [[k, mean_list(val)] for k, val in this_card_res_diss.items()]
          )
          this_card_res_obs = dict(
              [[k, mean_list(val)] for k, val in this_card_res_obs.items()]
          )
          for k, val in this_card_res_doc.items():
              if val == "none":  # no data collected on differentials of this size
                  continue
              else:  # if a value was collected for this value for doctors, it was collected for the other two algorithms too
                  if k not in doc_topn_caseav.keys():
                      doc_topn_caseav[k] = {
                          "count": 1,
                          "suff": deepcopy(this_card_res_suff[k]),
                          "diss": deepcopy(this_card_res_diss[k]),
                          "obs": deepcopy(this_card_res_obs[k]),
                          "doc": deepcopy(this_card_res_doc[k]),
                      }
                  else:  # we have recorded a case of this length before
                      doc_topn_caseav[k]["count"] += 1
                      doc_topn_caseav[k]["suff"] += deepcopy(this_card_res_suff[k])
                      doc_topn_caseav[k]["diss"] += deepcopy(this_card_res_diss[k])
                      doc_topn_caseav[k]["obs"] += deepcopy(this_card_res_obs[k])
                      doc_topn_caseav[k]["doc"] += deepcopy(this_card_res_doc[k])
  
          # card mean score
          if [val[2] for val in [_val for _val in doc_res_n if _val[1] != 0]] == []:
              continue
          doc_mean_score = np.mean(
              [val[2] for val in [_val for _val in doc_res_n if _val[1] != 0]]
          )
          obs_mean_score = np.mean(
              [
                  deepcopy(pred_obs[val[1] - 1])
                  for val in [_val for _val in doc_res_n if _val[1] != 0]
              ]
          )
          doc_score += [doc_mean_score]
          obs_score += [obs_mean_score]
  
      doc_score = []
      doc_error = []
      obs_score = []
      obs_error = []
      suff_score = []
      suff_error = []
      diss_score = []
      diss_error = []
  
      for k, val in doc_topn.items():
          n = val["count"]
          if n < 50:
              continue
  
          docp = sum(val["doctor"].values())[1] / n
          obsp = sum(val["obs"].values())[1] / n
          suffp = sum(val["sufficiency"].values())[1] / n
          dissp = sum(val["disablement"].values())[1] / n
          doc_score += [docp]
          doc_error += [np.sqrt(docp * (1 - docp) / n)]
          obs_score += [obsp]
          obs_error += [np.sqrt(obsp * (1 - obsp) / n)]
          suff_score += [suffp]
          suff_error += [np.sqrt(suffp * (1 - suffp) / n)]
          diss_score += [suffp]
          diss_error += [np.sqrt(dissp * (1 - dissp) / n)]
  
      raw_data = {
          "doc_score": doc_score,
          "doc_error": doc_error,
          "obs_score": obs_score,
          "obs_error": obs_error,
          "sufficiency_score": suff_score,
          "sufficiency_error": suff_error,
          "disablement_score": diss_score,
          "disablement_error": diss_error,
      }
      df_results = pd.DataFrame(
          raw_data,
          columns=[
              "doc_score",
              "doc_error",
              "obs_score",
              "obs_error",
              "sufficiency_score",
              "sufficiency_error",
              "disablement_score",
              "disablement_error",
          ],
      )
  
      df_results.to_pickle(args.results / "supp_table_3_df.p")
  
      return df_results, doc_topn
  
  
  def make_table_two(*, args, df_results, doc_topn):
      comf_thresh = 0.05
      res = np.zeros((4, 4))
      total_number = 0
      for k, val in doc_topn.items():
          if val["count"] < 50:
              continue
          d_doc = [
              [1 for i in range(_val[1])] + [0 for i in range(_val[0] - _val[1])]
              for _val in val["doctor"].values()
          ]
          fd_doc = [item for sublist in d_doc for item in sublist]
          d_obs = [
              [1 for i in range(_val[1])] + [0 for i in range(_val[0] - _val[1])]
              for _val in val["obs"].values()
          ]
          fd_obs = [item for sublist in d_obs for item in sublist]
          d_sufficiency = [
              [1 for i in range(_val[1])] + [0 for i in range(_val[0] - _val[1])]
              for _val in val["sufficiency"].values()
          ]
          fd_sufficiency = [item for sublist in d_sufficiency for item in sublist]
          d_disablement = [
              [1 for i in range(_val[1])] + [0 for i in range(_val[0] - _val[1])]
              for _val in val["disablement"].values()
          ]
          fd_disablement = [item for sublist in d_disablement for item in sublist]
  
          pvalue_matrix = np.array(
              [
                  [
                      0,
                      bintest(fd_doc, fd_obs, comf_thresh),
                      bintest(fd_doc, fd_sufficiency, comf_thresh),
                      bintest(fd_doc, fd_disablement, comf_thresh),
                  ],
                  [
                      bintest(fd_obs, fd_doc, comf_thresh),
                      0,
                      bintest(fd_obs, fd_sufficiency, comf_thresh),
                      bintest(fd_obs, fd_disablement, comf_thresh),
                  ],
                  [
                      bintest(fd_sufficiency, fd_doc, comf_thresh),
                      bintest(fd_sufficiency, fd_obs, comf_thresh),
                      0,
                      bintest(fd_sufficiency, fd_disablement, comf_thresh),
                  ],
                  [
                      bintest(fd_disablement, fd_doc, comf_thresh),
                      bintest(fd_disablement, fd_obs, comf_thresh),
                      bintest(fd_disablement, fd_sufficiency, comf_thresh),
                      0,
                  ],
              ]
          )
  
          res += pvalue_matrix
  
          total_number += 1
      print(np.around(res, decimals=3))
  
      print(
          f"Doctor Score:\t {round(df_results.doc_score.mean(), 4)}\n"
          f"Obs Score:\t{round(df_results.obs_score.mean(), 4)}\n"
          f"Suff Score: \t{round(df_results.sufficiency_score.mean(), 4)}\n"
          f"Disab Score:\t{round(df_results.disablement_score.mean(), 4)}\n"
      )
  
  
  def make_figure_four(*, args, df_results):
      cmap, norm = matplotlib.colors.from_levels_and_colors(
          [0, 0.499, 0.501, 1], ["blue", "red", "green"]
      )
  
      x = df_results.doc_score
      y = df_results.obs_score
      col = [(x[i] - y[i]) + 0.5 for i in range(len(x))]
      plt.figure(figsize=(6, 4))
      plt.scatter(x, y, c=col, cmap=cmap, norm=norm)
      xticks = np.linspace(min(min(x), min(y)), max(max(x), max(y)), 100)
      yticks = xticks
      plt.xticks(np.arange(0.5, 0.95, 0.1))
      plt.yticks(np.arange(0.5, 0.95, 0.1))
      plt.plot(xticks, yticks, "-r", linestyle="--")
  
      plt.savefig(args.results / "obs_vs_doc.pdf")
  
      plt.clf()
      x = df_results.doc_score
      y = df_results.sufficiency_score
      col = [(x[i] - y[i]) + 0.5 for i in range(len(x))]
      plt.scatter(x, y, c=col, cmap=cmap, norm=norm)
      xticks = np.linspace(min(min(x), min(y)), max(max(x), max(y)), 100)
      yticks = xticks
      plt.xticks(np.arange(0.5, 0.95, 0.1))
      plt.yticks(np.arange(0.5, 0.95, 0.1))
      plt.plot(xticks, yticks, "-r", linestyle="--")
  
      plt.savefig(args.results / "suff_vs_doc.pdf")
  