{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d622b051",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import glob\n",
    "import pandas as pd\n",
    "import pyarrow.parquet as pq\n",
    "from tqdm import tqdm\n",
    "\n",
    "ROOT = \"./droid_1.0.1_first_20_chunks\"\n",
    "NUM_CHUNKS = 20\n",
    "OUT_CSV = \"droid_first_20_chunks_markovian_split.csv\"\n",
    "\n",
    "MARKOVIAN_KEYWORDS = [\n",
    "    \"pick up\", \"pick\", \"grasp\", \"grab\", \"lift\",\n",
    "    \"place\", \"put\", \"move\", \"push\", \"pull\",\n",
    "    \"open\", \"close\", \"turn on\", \"turn off\",\n",
    "    \"press\", \"insert\", \"remove\", \"slide\",\n",
    "]\n",
    "\n",
    "NON_MARKOVIAN_KEYWORDS = [\n",
    "    \"then\", \"after\", \"before\", \"first\", \"second\", \"third\", \"finally\", \"next\",\n",
    "    \"all\", \"every\", \"each\", \"remaining\", \"another\",\n",
    "    \"sort\", \"arrange\", \"organize\", \"clean up\", \"cleanup\",\n",
    "    \"stack\", \"unstack\", \"set up\", \"prepare\",\n",
    "    \"repeat\", \"again\", \"until\",\n",
    "    \"in order\", \"sequence\",\n",
    "]\n",
    "\n",
    "\n",
    "def normalize_text(x):\n",
    "    if x is None:\n",
    "        return \"\"\n",
    "    return str(x).lower().strip()\n",
    "\n",
    "\n",
    "def classify_task(text):\n",
    "    text = normalize_text(text)\n",
    "\n",
    "    non_score = 0\n",
    "    markov_score = 0\n",
    "\n",
    "    for kw in NON_MARKOVIAN_KEYWORDS:\n",
    "        if re.search(rf\"\\b{re.escape(kw)}\\b\", text):\n",
    "            non_score += 2\n",
    "\n",
    "    for kw in MARKOVIAN_KEYWORDS:\n",
    "        if re.search(rf\"\\b{re.escape(kw)}\\b\", text):\n",
    "            markov_score += 1\n",
    "\n",
    "    # Multiple action verbs usually means longer-horizon behavior\n",
    "    action_count = sum(\n",
    "        1 for kw in MARKOVIAN_KEYWORDS\n",
    "        if re.search(rf\"\\b{re.escape(kw)}\\b\", text)\n",
    "    )\n",
    "\n",
    "    if action_count >= 2 and any(w in text for w in [\"then\", \"after\", \"before\", \"and\"]):\n",
    "        non_score += 2\n",
    "\n",
    "    # Broad household goals are often non-Markovian because progress matters\n",
    "    broad_goal_words = [\"clean\", \"tidy\", \"organize\", \"arrange\", \"sort\"]\n",
    "    if any(w in text for w in broad_goal_words):\n",
    "        non_score += 2\n",
    "\n",
    "    if non_score > markov_score:\n",
    "        return \"non_markovian\"\n",
    "\n",
    "    return \"markovian\"\n",
    "\n",
    "\n",
    "def read_episode_summary(parquet_path):\n",
    "    columns = [\n",
    "        \"episode_index\",\n",
    "        \"task_index\",\n",
    "        \"language_instruction\",\n",
    "        \"language_instruction_2\",\n",
    "        \"language_instruction_3\",\n",
    "        \"task_category\",\n",
    "        \"is_episode_successful\",\n",
    "    ]\n",
    "\n",
    "    table = pq.read_table(\n",
    "        parquet_path,\n",
    "        columns=[c for c in columns if c in pq.read_schema(parquet_path).names],\n",
    "    )\n",
    "\n",
    "    df = table.slice(0, 1).to_pandas()\n",
    "    row = df.iloc[0].to_dict()\n",
    "\n",
    "    instruction_parts = [\n",
    "        row.get(\"language_instruction\", \"\"),\n",
    "        row.get(\"language_instruction_2\", \"\"),\n",
    "        row.get(\"language_instruction_3\", \"\"),\n",
    "        row.get(\"task_category\", \"\"),\n",
    "    ]\n",
    "\n",
    "    instruction_text = \" | \".join(\n",
    "        [normalize_text(x) for x in instruction_parts if normalize_text(x)]\n",
    "    )\n",
    "\n",
    "    return {\n",
    "        \"episode_index\": row.get(\"episode_index\"),\n",
    "        \"task_index\": row.get(\"task_index\"),\n",
    "        \"language_instruction\": row.get(\"language_instruction\", \"\"),\n",
    "        \"language_instruction_2\": row.get(\"language_instruction_2\", \"\"),\n",
    "        \"language_instruction_3\": row.get(\"language_instruction_3\", \"\"),\n",
    "        \"task_category\": row.get(\"task_category\", \"\"),\n",
    "        \"is_episode_successful\": row.get(\"is_episode_successful\"),\n",
    "        \"instruction_text\": instruction_text,\n",
    "        \"label\": classify_task(instruction_text),\n",
    "        \"parquet_path\": parquet_path,\n",
    "    }\n",
    "\n",
    "\n",
    "all_files = []\n",
    "\n",
    "for i in range(NUM_CHUNKS):\n",
    "    chunk = f\"chunk-{i:03d}\"\n",
    "    pattern = os.path.join(ROOT, \"data\", chunk, \"*.parquet\")\n",
    "    all_files.extend(sorted(glob.glob(pattern)))\n",
    "\n",
    "print(f\"Found {len(all_files)} parquet files\")\n",
    "\n",
    "rows = []\n",
    "\n",
    "for path in tqdm(all_files):\n",
    "    try:\n",
    "        rows.append(read_episode_summary(path))\n",
    "    except Exception as e:\n",
    "        rows.append({\n",
    "            \"episode_index\": None,\n",
    "            \"task_index\": None,\n",
    "            \"instruction_text\": \"\",\n",
    "            \"label\": \"read_error\",\n",
    "            \"parquet_path\": path,\n",
    "            \"error\": str(e),\n",
    "        })\n",
    "\n",
    "df = pd.DataFrame(rows)\n",
    "\n",
    "df.to_csv(OUT_CSV, index=False)\n",
    "\n",
    "print(df[\"label\"].value_counts())\n",
    "print(f\"Saved to {OUT_CSV}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}