{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "from IPython.display import Video\n",
    "\n",
    "\n",
    "def get_label(ans):\n",
    "    try:\n",
    "        ans = int(ans)\n",
    "    except:\n",
    "        print(\"Wrong Input\")\n",
    "        return False\n",
    "    if ans not in [1,2,3]:\n",
    "        print(\"Invalid option.\")\n",
    "        return False\n",
    "    if ans == 1:\n",
    "        return [1, 0]\n",
    "    elif ans == 2:\n",
    "        return [0, 1]\n",
    "    else:\n",
    "        return [0.5, 0.5]\n",
    "\n",
    "\n",
    "def create_human_label(save_dir, env_name, num_query=1000, start_idx=None, width=1000, height=500):\n",
    "    video_path = os.path.join(save_dir, env_name)\n",
    "    os.makedirs(os.path.join(video_path, \"label\"), exist_ok=True)\n",
    "    print(\"START!\")\n",
    "    if start_idx:\n",
    "        assert start_idx > 0, \"you must input with video number (1, 2, 3, ...)\"\n",
    "        interval = range(start_idx - 1, num_query)\n",
    "    else:\n",
    "        interval = range(num_query)\n",
    "        \n",
    "    for i in interval:\n",
    "        label = False\n",
    "        while not label:\n",
    "            print(f\"\\nVideo {i + 1}\")\n",
    "            video_file = os.path.join(video_path, f\"idx{i}.mp4\")\n",
    "            display(Video(video_file, width=width, height=height, html_attributes=\"loop autoplay\"))\n",
    "            reward = input(f\"[{i + 1}/{num_query}] Put Preference (1 (left), 2 (right), 3 (equal)):  \").strip()\n",
    "            label = get_label(reward)\n",
    "            if label:\n",
    "                with open(os.path.join(video_path, \"label\", f\"label_{i}.txt\"), \"w\") as f:\n",
    "                    f.write(reward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### create human label in save_dir, you could fix the start point.\n",
    "create_human_label(save_dir=\"../video\", env_name=\"antmaze-large-diverse-v2\", start_idx=956, num_query=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import pickle\n",
    "import numpy as np\n",
    "from tqdm import trange\n",
    "\n",
    "# make final pickle file from separated label files.\n",
    "def merge_labels(save_dir, env_name=\"antmaze-medium-play-v2\", num_query=1000, query_len=100, seed=3407):\n",
    "    label_dir = os.path.join(save_dir, env_name, \"label\")\n",
    "    # label_files = sorted(glob.glob(os.path.join(label_dir, \"*.txt\")), key=lambda x: int(x.split(\".\")[0].split(\"_\")[-1]))\n",
    "    labels = []\n",
    "    for idx in trange(num_query):\n",
    "        assert os.path.exists(os.path.join(label_dir, f\"label_{idx}.txt\")), f\"labeling is not finished. {idx + 1} / {num_query}\"\n",
    "        with open(os.path.join(label_dir, f\"label_{idx}.txt\")) as f:\n",
    "            choice = int(f.read().strip())\n",
    "            if choice == 1:\n",
    "                _label = 0\n",
    "            elif choice == 2:\n",
    "                _label = 1\n",
    "            elif choice == 3:\n",
    "                _label = -1\n",
    "        labels.append(_label)\n",
    "        \n",
    "    # labels = np.array(labels)\n",
    "        \n",
    "    with open(os.path.join(save_dir, env_name, f\"human_labels_numq{num_query}_len{query_len}_s{seed}.pkl\"), \"wb\") as f:\n",
    "        pickle.dump(labels, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:00<00:00, 22433.63it/s]\n"
     ]
    }
   ],
   "source": [
    "merge_labels(save_dir=\"../video\", env_name=\"antmaze-medium-play-v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:00<00:00, 19003.26it/s]\n"
     ]
    }
   ],
   "source": [
    "merge_labels(save_dir=\"../video\", env_name=\"antmaze-medium-diverse-v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:00<00:00, 20405.77it/s]\n"
     ]
    }
   ],
   "source": [
    "merge_labels(save_dir=\"../video\", env_name=\"antmaze-large-diverse-v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"../video/antmaze-medium-play-v2/human_labels_numq1000_len100_s3407.pkl\", \"rb\") as f:\n",
    "    labels = pickle.load(f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6.8 64-bit",
   "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.6.8"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
