{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
    "import numpy as np\n",
    "import sklearn\n",
    "import sklearn.metrics as skm\n",
    "import sklearn.ensemble as ensemble\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import trange, tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data_df.shape (3415, 23)\n",
      "data_df.columns Index(['src', 'mt', 'ref', 'comet', 'seqlogprob', 'attn_ign_src', 'laser',\n",
      "       'chrf2', 'alti', 'cometkiwi', 'labse', 'wass_to_unif', 'wass_to_data',\n",
      "       'wass_combo', 'is_hall', 'is_osc', 'is_fd', 'is_sd', 'attn', 'src_ids',\n",
      "       'mt_ids', 'src_len', 'mt_len'],\n",
      "      dtype='object')\n",
      "ref_df.shape (99531, 20)\n",
      "ref_df.columns Index(['idx', 'src', 'mt', 'ref', 'src_ids', 'mt_ids', 'ref_ids', 'seqlogprob',\n",
      "       'attn', 'comet', 'comet-qe', 'laser', 'unnatended_src_prop', 'alti',\n",
      "       'wass_to_data', 'wass_to_unif', 'labse', 'cometkiwi', 'mt_len',\n",
      "       'wass_combo'],\n",
      "      dtype='object')\n"
     ]
    }
   ],
   "source": [
    "file = 'ot-hallucination-detection-data.zip/hallucinations_deen_w_stats_and_scores.pkl'\n",
    "ref = 'ot-hallucination-detection-data.zip/ref_with_wass_combo.pkl'\n",
    "data = np.load(file, allow_pickle=True)\n",
    "data_df = pd.DataFrame(data)\n",
    "ref_df = pd.DataFrame(np.load(ref, allow_pickle=True))\n",
    "ref_df.rename(columns={'score': 'seqlogprob'}, inplace=True)\n",
    "all_metrics = ['comet', 'seqlogprob', 'attn_ign_src', 'laser', 'chrf2', 'alti', 'cometkiwi', 'labse', 'wass_to_unif', 'wass_to_data', 'wass_combo']\n",
    "metrics_to_inverse = ['attn_ign_src', 'wass_to_unif', 'wass_to_data', 'wass_combo']\n",
    "for metric in metrics_to_inverse:\n",
    "    data_df[metric] = - data_df[metric]\n",
    "\n",
    "# print info about testdata\n",
    "print('data_df.shape', data_df.shape)\n",
    "print('data_df.columns', data_df.columns)\n",
    "# print info about ref\n",
    "print('ref_df.shape', ref_df.shape)\n",
    "print('ref_df.columns', ref_df.columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hallucination Detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def auroc_fpr(pred, num_thresh = 1000):\n",
    "    thrs = np.linspace(start=pred.min(), stop=pred.max(), num=num_thresh)\n",
    "    decision_by_thr = np.zeros((len(pred), num_thresh))\n",
    "    \n",
    "    for i in range(len(thrs)):\n",
    "        thr = thrs[i]\n",
    "        y_pred = np.where(pred <= thr, 1, 0)\n",
    "        decision_by_thr[:, i] = y_pred\n",
    "    return decision_by_thr\n",
    "\n",
    "def calculate_auc_fp_from_pred(pred, testdata_is_hall):\n",
    "    decision_by_thr = auroc_fpr(pred, num_thresh = 1000)\n",
    "    is_hall = np.zeros(decision_by_thr.shape)\n",
    "    for i in range(is_hall.shape[1]):\n",
    "        is_hall[:, i] = testdata_is_hall.to_numpy()\n",
    "\n",
    "    tp = np.where((decision_by_thr == 1) & (is_hall==1), 1, 0)\n",
    "    fn = np.where((decision_by_thr == 0) & (is_hall==1), 1, 0)\n",
    "    tn = np.where((decision_by_thr == 0) & (is_hall==0), 1, 0)\n",
    "    fp = np.where((decision_by_thr == 1) & (is_hall==0), 1, 0)\n",
    "\n",
    "    tpr = tp.sum(axis=0) / (tp.sum(axis=0) + fn.sum(axis=0))\n",
    "    fpr = fp.sum(axis=0) / (fp.sum(axis=0) + tn.sum(axis=0))\n",
    "\n",
    "    x = np.interp(0.90, np.sort(tpr), np.sort(fpr))\n",
    "    auc, fp = skm.auc(np.sort(fpr), np.sort(tpr)) * 100, x * 100\n",
    "    return auc, fp\n",
    "\n",
    "def calculate_auc_fp(trainset_data_detectors,testset_data_detectors,testdata_is_hall,prediction_method,num_try = 10,*args, **kwargs):\n",
    "    all_auc, all_fp = 0, 0\n",
    "    for _ in range(num_try):\n",
    "        pred = prediction_method(trainset_data_detectors,testset_data_detectors,*args, **kwargs)\n",
    "        auc, fp = calculate_auc_fp_from_pred(pred, testdata_is_hall)\n",
    "        all_auc += auc\n",
    "        all_fp += fp\n",
    "    return all_auc / num_try, all_fp / num_try, pred\n",
    "\n",
    "def get_best_threshold(pred, testdata_is_hall,num_thresh = 1000):\n",
    "    thrs = np.linspace(start=pred.min(), stop=pred.max(), num=num_thresh)\n",
    "    decision_by_thr = auroc_fpr(pred, num_thresh = num_thresh)\n",
    "    is_hall = np.zeros(decision_by_thr.shape)\n",
    "    for i in range(is_hall.shape[1]):\n",
    "        is_hall[:, i] = testdata_is_hall.to_numpy()\n",
    "    f1 = np.zeros(decision_by_thr.shape[1])\n",
    "    for i in range(is_hall.shape[1]):\n",
    "        f1[i] = skm.f1_score(is_hall[:, i], decision_by_thr[:, i])\n",
    "    return thrs[np.argmax(f1)], f1.max()\n",
    "\n",
    "def plot_prediction(pred, testdata_is_hall, title, thr, nbins=50,legend=True):\n",
    "    plt.figure(figsize=(10, 3))\n",
    "    bins = np.histogram(pred, bins=nbins)[1]\n",
    "\n",
    "    sns.histplot(pred[testdata_is_hall == 0], bins=bins, color='green', label=\"No Hallucination\", kde=True, edgecolor='white', stat='percent')\n",
    "    sns.histplot(pred[testdata_is_hall == 1], bins=bins, color='red', label=\"Hallucination\", kde=True, edgecolor='white', stat='percent')\n",
    "    # plt.axvline(x=thr, color='black', linestyle='--', linewidth=2, label='Threshold')\n",
    "\n",
    "    plt.xlabel('Score', fontsize=30)\n",
    "    plt.ylabel('Proportion', fontsize=30)\n",
    "    # plt.title(title, fontsize=16)\n",
    "    sns.set_palette(\"Set2\")\n",
    "    if legend:\n",
    "        plt.legend(loc='upper left', fontsize=30)\n",
    "    plt.grid(True, linestyle='--', alpha=0.7)\n",
    "    sns.despine()\n",
    "    plt.xticks(fontsize=25)\n",
    "    plt.yticks(fontsize=25)\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prediction methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prediction_method_single(trainset_data_detectors,testset_data_detectors): # single detector\n",
    "    pred = testset_data_detectors\n",
    "    return pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prediction_method_IF(trainset_data_detectors,testset_data_detectors): # aggregation of detectors using isolation forest\n",
    "    isolation_forest = ensemble.IsolationForest(n_estimators=500,n_jobs=-1)\n",
    "    clf = isolation_forest.fit(trainset_data_detectors)\n",
    "    pred = clf.score_samples(testset_data_detectors)\n",
    "    return pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression\n",
    "def prediction_logistic_regression(trainset_data_detectors,testset_data_detectors,trainset_data_is_hall):\n",
    "    logistic_regression = LogisticRegression(max_iter=1000,solver='liblinear')\n",
    "    trainset_data_is_hall = np.array(trainset_data_is_hall).ravel()\n",
    "    logistic_regression.fit(trainset_data_detectors, trainset_data_is_hall)\n",
    "    pred = logistic_regression.predict_proba(testset_data_detectors)[:, 1]\n",
    "    return -pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "oods = ['is_hall', 'is_osc', 'is_fd', 'is_sd']\n",
    "\n",
    "data_df_trainset = ref_df\n",
    "data_df_testset = data_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_metrics = [\"alti\",\"seqlogprob\",\"wass_combo\"]\n",
    "external_metrics = [\"cometkiwi\",\"labse\"]\n",
    "model_external_metrics = [\"alti\",\"seqlogprob\",\"wass_combo\",\"cometkiwi\",\"labse\"]\n",
    "\n",
    "metrics_type = {\n",
    "    \"model\":model_metrics,\n",
    "    \"external\":external_metrics,\n",
    "    \"all\":model_external_metrics}\n",
    "sorted_metrics = [\"labse\",\"wass_combo\",\"cometkiwi\",\"alti\",\"seqlogprob\"]\n",
    "\n",
    "\n",
    "normalized_metrics = {}\n",
    "normalized_metrics_trainset = {}\n",
    "for metric in model_external_metrics:\n",
    "    minnn = data_df_trainset[metric].min()\n",
    "    maxxx = data_df_trainset[metric].max()\n",
    "    # Min-max feature scaling on the trainset \n",
    "    normalized_metrics_trainset[metric] = data_df_trainset[metric].apply(lambda x: (x - minnn) / (maxxx - minnn))\n",
    "    normalized_metrics[metric] = data_df_testset[metric].apply(lambda x: (x - minnn) / (maxxx - minnn))\n",
    "\n",
    "\n",
    "# ablation study\n",
    "sizes = [2,5,10,50,100,500,1000,5000,10000,50000,len(data_df_trainset)]\n",
    "Nd = 10 # number of datasets per size\n",
    "datasets = {}\n",
    "for size in sizes:\n",
    "    datasets[size] = {}\n",
    "    for i in range(Nd):\n",
    "        datasets[size][i] = data_df_trainset.sample(size,random_state=size+i)\n",
    "\n",
    "normalized_metrics_datasets_trainset = {}\n",
    "normalized_metrics_testset = {}\n",
    "for size in sizes:\n",
    "    normalized_metrics_datasets_trainset[size] = {}\n",
    "    normalized_metrics_testset[size] = {}\n",
    "    for i in range(Nd):\n",
    "        normalized_metrics_datasets_trainset[size][i] = {}\n",
    "        normalized_metrics_testset[size][i] = {}\n",
    "        for metric in model_external_metrics:\n",
    "            minnn = datasets[size][i][metric].min()\n",
    "            maxxx = datasets[size][i][metric].max()\n",
    "            # Min-max feature scaling on the trainset \n",
    "            normalized_metrics_datasets_trainset[size][i][metric] = datasets[size][i][metric].apply(lambda x: (x - minnn) / (maxxx - minnn))\n",
    "            normalized_metrics_testset[size][i][metric] = data_df_testset[metric].apply(lambda x: (x - minnn) / (maxxx - minnn))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "normalized_metrics = pd.DataFrame(normalized_metrics)\n",
    "normalized_metrics_trainset = pd.DataFrame(normalized_metrics_trainset)\n",
    "for size in sizes:\n",
    "    for i in range(Nd):\n",
    "        normalized_metrics_datasets_trainset[size][i] = pd.DataFrame(normalized_metrics_datasets_trainset[size][i])\n",
    "        normalized_metrics_testset[size][i] = pd.DataFrame(normalized_metrics_testset[size][i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "configs = []\n",
    "# ablation study\n",
    "# for ood in [\"is_hall\"]:\n",
    "#     for size in sizes:\n",
    "#         for i in range(Nd):\n",
    "#             # # ablation Sum_norm\n",
    "#             config = {\"ood\":ood,\"name\":\"Sum_norm_\"+str(size)+\"_\"+str(i),\"trainset_data_detectors\":None,\n",
    "#                     \"testset_data_detectors\":np.sum(normalized_metrics_testset[size][i][model_external_metrics],axis=1),\n",
    "#                     \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_method_single,\"num_try\":1,\n",
    "#                     \"trainset_data_is_hall\":None, \"method\":None}\n",
    "#             configs.append(config)\n",
    "#             # # ablation Max_norm\n",
    "#             config = {\"ood\":ood,\"name\":\"Max_norm_\"+str(size)+\"_\"+str(i),\"trainset_data_detectors\":None,\n",
    "#                     \"testset_data_detectors\":np.max(normalized_metrics_testset[size][i][model_external_metrics],axis=1),\n",
    "#                     \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_method_single,\"num_try\":1,\n",
    "#                     \"trainset_data_is_hall\":None, \"method\":None}\n",
    "#             configs.append(config)\n",
    "#             # # ablation IF\n",
    "#             config = {\"ood\":ood,\"name\":\"IF_\"+str(size)+\"_\"+str(i),\"trainset_data_detectors\":normalized_metrics_datasets_trainset[size][i][model_external_metrics],\n",
    "#                     \"testset_data_detectors\":normalized_metrics_testset[size][i][model_external_metrics],\n",
    "#                     \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_method_IF,\"num_try\":3,\n",
    "#                     \"trainset_data_is_hall\":None, \"method\":None}\n",
    "#             configs.append(config)\n",
    "\n",
    "# performance of the different methods\n",
    "for ood in oods:\n",
    "# for ood in [\"is_hall\"]:\n",
    "    for metric in model_external_metrics:\n",
    "        config = {\"ood\":ood,\"name\":metric,\"trainset_data_detectors\":None,\"testset_data_detectors\":data_df_testset[metric],\n",
    "                \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_method_single,\"num_try\":1,\n",
    "                \"trainset_data_is_hall\":None, \"method\":None}\n",
    "        configs.append(config)\n",
    "    # Sum_norm\n",
    "    for name,metrics in metrics_type.items():\n",
    "        config = {\"ood\":ood,\"name\":\"Sum_norm_\"+name,\"trainset_data_detectors\":None,\n",
    "                \"testset_data_detectors\":np.sum(pd.DataFrame(normalized_metrics)[metrics],axis=1),\n",
    "                \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_method_single,\"num_try\":1,\n",
    "                \"trainset_data_is_hall\":None, \"method\":None}\n",
    "        configs.append(config)\n",
    "    # Max_norm\n",
    "    for name,metrics in metrics_type.items():\n",
    "        config = {\"ood\":ood,\"name\":\"Max_norm_\"+name,\"trainset_data_detectors\":None,\n",
    "                \"testset_data_detectors\":np.max(pd.DataFrame(normalized_metrics)[metrics],axis=1),\n",
    "                \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_method_single,\"num_try\":1,\n",
    "                \"trainset_data_is_hall\":None, \"method\":None}\n",
    "        configs.append(config)\n",
    "    # IF\n",
    "    for name,metrics in metrics_type.items():\n",
    "        config = {\"ood\":ood,\"name\":\"IF_\"+name,\"trainset_data_detectors\":data_df_trainset[metrics],\n",
    "                \"testset_data_detectors\":data_df_testset[metrics],\n",
    "                \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_method_IF,\"num_try\":10,\n",
    "                \"trainset_data_is_hall\":None, \"method\":None}\n",
    "        configs.append(config)\n",
    "\n",
    "    # logistic_regression\n",
    "    # config = {\"ood\":ood,\"name\":\"logistic_regression\",\"trainset_data_detectors\":data_df_testset[model_external_metrics],\n",
    "    #             \"testset_data_detectors\":data_df_testset[model_external_metrics],\n",
    "    #             \"testdata_is_hall\":data_df_testset[ood],\"prediction_method\":prediction_logistic_regression,\"num_try\":1,\n",
    "    #             \"trainset_data_is_hall\":data_df_testset[ood]}\n",
    "    # configs.append(config)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_full = {}\n",
    "for config in configs:\n",
    "    if config[\"ood\"] not in results_full:\n",
    "        results_full[config[\"ood\"]] = {}\n",
    "    results_full[config[\"ood\"]][config[\"name\"]] = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for config in tqdm(configs):\n",
    "    if config[\"name\"] not in results_full[config[\"ood\"]] or len(results_full[config[\"ood\"]][config[\"name\"]]) < 1:\n",
    "        print(config[\"ood\"], config[\"name\"])\n",
    "        if config[\"trainset_data_is_hall\"] is None:\n",
    "            if config[\"method\"] is None:\n",
    "                auc, fp, pred = calculate_auc_fp(config[\"trainset_data_detectors\"],config[\"testset_data_detectors\"], config[\"testdata_is_hall\"], \n",
    "                                            config[\"prediction_method\"], num_try = config[\"num_try\"])\n",
    "            else:\n",
    "                auc, fp, pred = calculate_auc_fp(config[\"trainset_data_detectors\"],config[\"testset_data_detectors\"], config[\"testdata_is_hall\"], \n",
    "                                            config[\"prediction_method\"], num_try = config[\"num_try\"], method = config[\"method\"])\n",
    "        else:\n",
    "            auc, fp, pred = calculate_auc_fp(config[\"trainset_data_detectors\"],config[\"testset_data_detectors\"], config[\"testdata_is_hall\"], \n",
    "                                         config[\"prediction_method\"], num_try = config[\"num_try\"],\n",
    "                                         trainset_data_is_hall =config[\"trainset_data_is_hall\"])\n",
    "        thr, f1 = get_best_threshold(pred, config[\"testdata_is_hall\"])\n",
    "        results_full[config[\"ood\"]][config[\"name\"]].append({\"auc\":auc,\"fp\":fp, \"f1_score\":f1, \"pred\":pred, \"thr\":thr})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save results_full\n",
    "# import pickle\n",
    "# with open('ot-hallucination-detection-data_results_full.pkl', 'wb') as f:\n",
    "#     pickle.dump(results_full, f)\n",
    "# import pickle\n",
    "# with open('results_ablation.pkl', 'wb') as f:\n",
    "#     pickle.dump(results_full, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # load results_full\n",
    "import pickle\n",
    "with open('ot-hallucination-detection-data_results_full.pkl', 'rb') as f:\n",
    "    results_full = pickle.load(f)\n",
    "# load ablation results\n",
    "# import pickle\n",
    "# with open('results_ablation.pkl', 'rb') as f:\n",
    "#     results_full = pickle.load(f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot ablation study\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# prepare seaborn lineplot\n",
    "sum_norm_df = pd.DataFrame(columns=[\"name\",\"size\",\"i\",\"auc\",\"fp\",\"f1_score\"]) #is_hall\n",
    "max_norm_df = pd.DataFrame(columns=[\"name\",\"size\",\"i\",\"auc\",\"fp\",\"f1_score\"]) #is_hall\n",
    "IF_df = pd.DataFrame(columns=[\"name\",\"size\",\"i\",\"auc\",\"fp\",\"f1_score\"]) #is_hall\n",
    "for ood in [\"is_hall\"]:\n",
    "    for size in sizes[1:]:\n",
    "        for i in range(Nd):\n",
    "            sum_norm_df = sum_norm_df.append({\"name\":\"Sum_norm_\"+str(size)+\"_\"+str(i),\"size\":size,\"i\":i,\"auc\":results_full[ood][\"Sum_norm_\"+str(size)+\"_\"+str(i)][0][\"auc\"],\"fp\":results_full[ood][\"Sum_norm_\"+str(size)+\"_\"+str(i)][0][\"fp\"],\"f1_score\":results_full[ood][\"Sum_norm_\"+str(size)+\"_\"+str(i)][0][\"f1_score\"]},ignore_index=True)\n",
    "            max_norm_df = max_norm_df.append({\"name\":\"Max_norm_\"+str(size)+\"_\"+str(i),\"size\":size,\"i\":i,\"auc\":results_full[ood][\"Max_norm_\"+str(size)+\"_\"+str(i)][0][\"auc\"],\"fp\":results_full[ood][\"Max_norm_\"+str(size)+\"_\"+str(i)][0][\"fp\"],\"f1_score\":results_full[ood][\"Max_norm_\"+str(size)+\"_\"+str(i)][0][\"f1_score\"]},ignore_index=True)\n",
    "            IF_df = IF_df.append({\"name\":\"IF_\"+str(size)+\"_\"+str(i),\"size\":size,\"i\":i,\"auc\":results_full[ood][\"IF_\"+str(size)+\"_\"+str(i)][0][\"auc\"],\"fp\":results_full[ood][\"IF_\"+str(size)+\"_\"+str(i)][0][\"fp\"],\"f1_score\":results_full[ood][\"IF_\"+str(size)+\"_\"+str(i)][0][\"f1_score\"]},ignore_index=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 110280 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 107008 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "import matplotlib as mpl \n",
    "mpl.rcParams['axes.spines.right'] = False \n",
    "mpl.rcParams['axes.spines.top'] = False\n",
    "sns.set_palette(\"Set2\")\n",
    "\n",
    "# Plot mean and variance for AUROC\n",
    "plt.figure(figsize=(10, 5))\n",
    "sns.lineplot(data=sum_norm_df, x=\"size\", y=\"auc\", estimator=\"mean\", ci=\"sd\", label=\"STARE\", marker=\"o\", linewidth=2, markersize=15)\n",
    "sns.lineplot(data=max_norm_df, x=\"size\", y=\"auc\", estimator=\"mean\", ci=\"sd\", label=\"Max-norm\", marker=\"s\", linewidth=2, markersize=15)\n",
    "sns.lineplot(data=IF_df, x=\"size\", y=\"auc\", estimator=\"mean\", ci=\"sd\", label=\"IF\", marker=\"^\", linewidth=2, markersize=15)\n",
    "plt.xlabel('Calibration Set Size', fontsize=35)\n",
    "plt.xscale('log')\n",
    "plt.ylabel('AUROC', fontsize=35)\n",
    "plt.xticks(fontsize=28)\n",
    "plt.yticks(fontsize=28)\n",
    "plt.axhline(y=91.7, color='r', linestyle='--', label=\"LaBSE\", linewidth=3)\n",
    "plt.legend(title=None, fontsize=28, loc='best')\n",
    "\n",
    "\n",
    "# plt.grid(True, linestyle='--', alpha=0.7)\n",
    "plt.show()\n",
    "\n",
    "# Plot mean and variance for FPR@90TPR\n",
    "plt.figure(figsize=(10, 5))\n",
    "sns.lineplot(data=sum_norm_df, x=\"size\", y=\"fp\", estimator=\"mean\", ci=\"sd\", label=\"STARE\", marker=\"o\",markersize=15,linewidth=2)\n",
    "sns.lineplot(data=max_norm_df, x=\"size\", y=\"fp\", estimator=\"mean\", ci=\"sd\", label=\"Max-norm\", marker=\"s\",markersize=15,linewidth=2)\n",
    "sns.lineplot(data=IF_df, x=\"size\", y=\"fp\", estimator=\"mean\", ci=\"sd\", label=\"IF\", marker=\"^\",markersize=15,linewidth=2)\n",
    "plt.xlabel('Calibration Set Size', fontsize=35)\n",
    "plt.xscale('log')\n",
    "plt.ylabel('FPR@90TPR', fontsize=35)\n",
    "plt.xticks(fontsize=28)\n",
    "plt.yticks(fontsize=28)\n",
    "plt.axhline(y=26.9, color='r', linestyle='--', label=\"LaBSE\", linewidth=3)\n",
    "plt.legend(title=None, fontsize=28, loc='best')\n",
    "\n",
    "\n",
    "# plt.grid(True, linestyle='--', alpha=0.7)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 78572 characters]",
      "text/plain": [
       "<Figure size 1000x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 67524 characters]",
      "text/plain": [
       "<Figure size 1000x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "m_to_plot = \"labse\"\n",
    "ood = \"is_osc\"\n",
    "i = 0\n",
    "\n",
    "# Filter the data where is_osc is False\n",
    "# is_osc_false = data_df_testset['is_osc'] == 0\n",
    "is_osc_false = data_df_testset['is_osc'] <=1\n",
    "# Plot the prediction for is_osc=False points\n",
    "plot_prediction(results_full[ood][m_to_plot][i][\"pred\"][is_osc_false], data_df_testset[ood][is_osc_false].values, \"\", results_full[ood][m_to_plot][i][\"thr\"], nbins=20,legend=True)\n",
    "\n",
    "m_to_plot = \"Sum_norm_all\"\n",
    "\n",
    "# Plot the prediction for is_osc=False points with modified values\n",
    "plot_prediction(results_full[ood][m_to_plot][i][\"pred\"][is_osc_false] / 5, data_df_testset[ood][is_osc_false].values, \"\", results_full[ood][m_to_plot][i][\"thr\"] / 5, nbins=20,legend=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from matplotlib.patches import Circle, RegularPolygon\n",
    "from matplotlib.path import Path\n",
    "from matplotlib.projections import register_projection\n",
    "from matplotlib.projections.polar import PolarAxes\n",
    "from matplotlib.spines import Spine\n",
    "from matplotlib.transforms import Affine2D\n",
    "\n",
    "\n",
    "def radar_factory(num_vars, frame='circle'):\n",
    "    \"\"\"\n",
    "    Create a radar chart with `num_vars` axes.\n",
    "\n",
    "    This function creates a RadarAxes projection and registers it.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    num_vars : int\n",
    "        Number of variables for radar chart.\n",
    "    frame : {'circle', 'polygon'}\n",
    "        Shape of frame surrounding axes.\n",
    "\n",
    "    \"\"\"\n",
    "    # calculate evenly-spaced axis angles\n",
    "    theta = np.linspace(0, 2*np.pi, num_vars, endpoint=False)\n",
    "\n",
    "    class RadarTransform(PolarAxes.PolarTransform):\n",
    "\n",
    "        def transform_path_non_affine(self, path):\n",
    "            # Paths with non-unit interpolation steps correspond to gridlines,\n",
    "            # in which case we force interpolation (to defeat PolarTransform's\n",
    "            # autoconversion to circular arcs).\n",
    "            if path._interpolation_steps > 1:\n",
    "                path = path.interpolated(num_vars)\n",
    "            return Path(self.transform(path.vertices), path.codes)\n",
    "\n",
    "    class RadarAxes(PolarAxes):\n",
    "\n",
    "        name = 'radar'\n",
    "        PolarTransform = RadarTransform\n",
    "\n",
    "        def __init__(self, *args, **kwargs):\n",
    "            super().__init__(*args, **kwargs)\n",
    "            # rotate plot such that the first axis is at the top\n",
    "            self.set_theta_zero_location('N')\n",
    "\n",
    "        def fill(self, *args, closed=True, **kwargs):\n",
    "            \"\"\"Override fill so that line is closed by default\"\"\"\n",
    "            return super().fill(closed=closed, *args, **kwargs)\n",
    "\n",
    "        def plot(self, *args, **kwargs):\n",
    "            \"\"\"Override plot so that line is closed by default\"\"\"\n",
    "            lines = super().plot(*args, **kwargs)\n",
    "            for line in lines:\n",
    "                self._close_line(line)\n",
    "\n",
    "        def _close_line(self, line):\n",
    "            x, y = line.get_data()\n",
    "            # FIXME: markers at x[0], y[0] get doubled-up\n",
    "            if x[0] != x[-1]:\n",
    "                x = np.append(x, x[0])\n",
    "                y = np.append(y, y[0])\n",
    "                line.set_data(x, y)\n",
    "\n",
    "        def set_varlabels(self, labels):\n",
    "            self.set_thetagrids(np.degrees(theta), labels)\n",
    "\n",
    "        def _gen_axes_patch(self):\n",
    "            # The Axes patch must be centered at (0.5, 0.5) and of radius 0.5\n",
    "            # in axes coordinates.\n",
    "            if frame == 'circle':\n",
    "                return Circle((0.5, 0.5), 0.5)\n",
    "            elif frame == 'polygon':\n",
    "                return RegularPolygon((0.5, 0.5), num_vars,\n",
    "                                      radius=.5, edgecolor=\"k\")\n",
    "            else:\n",
    "                raise ValueError(\"Unknown value for 'frame': %s\" % frame)\n",
    "\n",
    "        def _gen_axes_spines(self):\n",
    "            if frame == 'circle':\n",
    "                return super()._gen_axes_spines()\n",
    "            elif frame == 'polygon':\n",
    "                # spine_type must be 'left'/'right'/'top'/'bottom'/'circle'.\n",
    "                spine = Spine(axes=self,\n",
    "                              spine_type='circle',\n",
    "                              path=Path.unit_regular_polygon(num_vars))\n",
    "                # unit_regular_polygon gives a polygon of radius 1 centered at\n",
    "                # (0, 0) but we want a polygon of radius 0.5 centered at (0.5,\n",
    "                # 0.5) in axes coordinates.\n",
    "                spine.set_transform(Affine2D().scale(.5).translate(.5, .5)\n",
    "                                    + self.transAxes)\n",
    "                return {'polar': spine}\n",
    "            else:\n",
    "                raise ValueError(\"Unknown value for 'frame': %s\" % frame)\n",
    "\n",
    "    register_projection(RadarAxes)\n",
    "    return theta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 116268 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from matplotlib import cm\n",
    "N = len(oods)\n",
    "theta = radar_factory(N, frame='polygon')\n",
    "spoke_labels = oods\n",
    "\n",
    "method_name_mapping = {\n",
    "    'alti':\"ALTI+ (Ferrando, et al. 2022)\",\n",
    "    \"seqlogprob\":\"Seq-Logprob (Guerreiro et al. 2022)\",\n",
    "    \"wass_combo\":\"Wass-Combo (Guerreiro et al. 2022)\",\n",
    "    \"cometkiwi\":\"CometKiwi (Rei et al. 2022)\",\n",
    "    \"labse\":\"LaBSE (Feng et al. 2020)\",\n",
    "    \"Sum_norm_all\":\"STARE (Our method)\"\n",
    "}\n",
    "\n",
    "dataset_to_label = {\n",
    "    \"is_hall\":\"Hallucination\",\n",
    "    \"is_osc\":\"Oscillation\",\n",
    "    \"is_fd\":\"Fully\\nDetached\",\n",
    "    \"is_sd\":\"Strongly\\nDetached\"\n",
    "}\n",
    "\n",
    "\n",
    "\n",
    "fig, axs = plt.subplots(figsize=(20, 5), subplot_kw=dict(projection='radar'))\n",
    "\n",
    "num_colors = len(method_name_mapping)\n",
    "# colors = cm.get_cmap('tab20', num_colors)\n",
    "colors = ['m', 'g', 'c', 'k', 'r', 'b']\n",
    "data = []\n",
    "for ood in oods:\n",
    "    data.append([results_full[ood][metric][0][\"auc\"] for metric in method_name_mapping])\n",
    "\n",
    "# Transpose data\n",
    "data = np.array(data).T\n",
    "def zooming_function(x):\n",
    "    # Define the power factor (you can adjust this based on the desired zooming effect)\n",
    "    power_factor = 5.0\n",
    "    \n",
    "    # Apply the power function to zoom in on the higher range (50 to 100)\n",
    "    return x**power_factor\n",
    "data_zoomed = zooming_function(data)\n",
    "\n",
    "# Plot\n",
    "# axs.set_title(\"Logarithmic Radar Plot\", weight='bold', size='medium', position=(0.5, 1.1),\n",
    "#               horizontalalignment='center', verticalalignment='center')\n",
    "\n",
    "for i, d in enumerate(data_zoomed):\n",
    "    # color = colors(i / num_colors)  # Get color from colormap\n",
    "    color = colors[i]\n",
    "    axs.plot(theta, d, color=color)\n",
    "    axs.fill(theta, d, facecolor=color, alpha=0.1, label='_nolegend_')\n",
    "axs.set_varlabels([dataset_to_label[ood] for ood in spoke_labels])\n",
    "for label in axs.get_xticklabels():\n",
    "    label.set_fontsize('large')  # Adjust font size\n",
    "    label.set_weight('bold')  # Add bold style\n",
    "    label.set_color('black')  # Change font color\n",
    "    label.set_backgroundcolor('white')  # Add white background to labels\n",
    "\n",
    "# Add legend (metric names)\n",
    "labels = [method_name_mapping[metric] for metric in method_name_mapping]\n",
    "legend = axs.legend(labels, loc='center left', bbox_to_anchor=(0.7, 1), labelspacing=0.1, fontsize='large')  # Move legend outside to the left\n",
    "# fig.text(0.5, 0.965, 'Logarithmic 5-Factor Solution Profiles Across Four Scenarios',\n",
    "#          horizontalalignment='center', color='black', weight='bold', size='large')\n",
    "axs.set_rgrids([zooming_function(val) for val in [50,70,80,90,95]], labels=['50','70','80','90','95'], angle=0)\n",
    "for label in axs.get_yticklabels():\n",
    "    label.set_fontsize('large')\n",
    "axs.set_axisbelow(False)  # <-- added\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
