{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ralstonia Annotation Tool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from io import StringIO\n",
    "import json\n",
    "from datetime import datetime as dt\n",
    "\n",
    "import pandas as pd\n",
    "import panel as pn\n",
    "\n",
    "import scripts.tap_const as tc\n",
    "import scripts.tap_image as ti"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pn.extension(\"tabulator\", design=\"bootstrap\")\n",
    "\n",
    "template = pn.template.BootstrapTemplate(title=\"Ralstonia Annotation Tool\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.set_option(\"display.max_colwidth\", 500)\n",
    "pd.set_option(\"display.max_columns\", 500)\n",
    "pd.set_option(\"display.width\", 1000)\n",
    "pd.set_option(\"display.max_rows\", 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EXPERIMENT = \"72AC_PhD_2404\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup Paths\n",
    "if tc.phenopsis.joinpath(EXPERIMENT).is_dir() is True:\n",
    "    pt_data = tc.data\n",
    "    pt_images = tc.phenopsis.joinpath(EXPERIMENT)\n",
    "    pt_rotations = tc.dataout.joinpath(\"rotation_angles\").joinpath(f\"{EXPERIMENT}\")\n",
    "else:\n",
    "    here = Path(\".\").parent\n",
    "    pt_data = here.joinpath(\"data\")\n",
    "    pt_images = here.joinpath(\"images\").joinpath(EXPERIMENT)\n",
    "    pt_rotations = here.joinpath(\"rotation_angles\").joinpath(f\"{EXPERIMENT}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## User Interface"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Source Selection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Download Template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dwn_template = pn.widgets.FileDownload(\n",
    "    file=pt_data.joinpath(f\"{EXPERIMENT}_raw.csv\"),\n",
    "    filename=f\"{EXPERIMENT}_raw.csv\",\n",
    "    button_type=\"success\",\n",
    "    label=\"Download template annotation file\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Upload existing file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_input = pn.widgets.FileInput(accept=\".csv,.json\")\n",
    "table = pn.widgets.Tabulator(\n",
    "    value=pd.DataFrame(),\n",
    "    pagination=\"local\",\n",
    "    page_size=20,\n",
    "    sizing_mode=\"stretch_width\",\n",
    ")\n",
    "\n",
    "ck_show_finished = pn.widgets.ToggleIcon(\n",
    "    icon=\"eye-x\", size=\"4em\", active_icon=\"eye\", value=False\n",
    ")\n",
    "mk_show_finished = pn.pane.Markdown(\"Hide complete plants\")\n",
    "\n",
    "@pn.depends(file_input.param.value, watch=True)\n",
    "def on_file_loaded(value):\n",
    "    if not value or not isinstance(value, bytes):\n",
    "        return pd.DataFrame()\n",
    "\n",
    "    string_io = StringIO(value.decode(\"utf8\"))\n",
    "    ret = pd.read_csv(string_io, sep=\";\").assign(\n",
    "        file_name=lambda x: x.filepath.str.replace(\".tif\", \".jpg\")\n",
    "    )\n",
    "    ret = ret[ret.treatment == \"RS\"]\n",
    "    if \"di\" not in ret:\n",
    "        ret[\"di\"] = 0\n",
    "    if \"done\" not in ret:\n",
    "        ret[\"done\"] = False\n",
    "\n",
    "    table.value = ret"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Annotation Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plant Selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sl_plant = pn.widgets.Select(name=\"Plant\", options=[])\n",
    "\n",
    "@pn.depends(table.param.value, ck_show_finished.param.value, watch=True)\n",
    "def on_table_changed(file: str, show_done: bool):\n",
    "    df = table.value\n",
    "    if \"done\" in df:\n",
    "        sl_plant.options = list(\n",
    "            df.plant.unique() if show_done is True else df[df.done == False].plant.unique()\n",
    "        )\n",
    "        mk_show_finished.object = (\n",
    "            \"Show complete plants\" if show_done is True else \"Hide complete plants\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Main Annotation UI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im_current = pn.pane.Image(max_width=800, max_height=800, sizing_mode=\"stretch_width\")\n",
    "discrete_player = pn.widgets.DiscretePlayer(\n",
    "    name=\"Discrete Player\", options=[0], value=0, loop_policy=\"loop\"\n",
    ")\n",
    "discrete_player.interval = 1000\n",
    "ii_disease_index = pn.widgets.IntInput(\n",
    "    name=\"Disease Index\",\n",
    "    start=0,\n",
    "    end=4,\n",
    "    step=1,\n",
    "    value=0,\n",
    "    max_width=80,\n",
    "    sizing_mode=\"stretch_width\",\n",
    ")\n",
    "\n",
    "dwn_annotations = pn.widgets.FileDownload(\n",
    "    file=pt_data.joinpath(f\"{EXPERIMENT}_raw.csv\"),\n",
    "    filename=f\"{EXPERIMENT}_raw.csv\",\n",
    "    label=\" \",\n",
    "    name=\"Download Annotations\",\n",
    "    sizing_mode=\"stretch_width\",\n",
    "    icon=\"file-download\",\n",
    "    button_type=\"primary\",\n",
    ")\n",
    "\n",
    "updating: bool = False\n",
    "\n",
    "\n",
    "@pn.cache(max_items=10, policy=\"LRU\")\n",
    "def get_plant_data(df, plant_name) -> pd.DataFrame:\n",
    "    file_names, rotations = [], []\n",
    "    for k, v in json.load(\n",
    "        open(pt_rotations.joinpath(sl_plant.value).with_suffix(\".json\"), \"r\")\n",
    "    ).items():\n",
    "        file_names.append(k + \".jpg\")\n",
    "        rotations.append(v)\n",
    "\n",
    "    return df[df.plant == plant_name].merge(\n",
    "        pd.DataFrame(data={\"file_name\": file_names, \"rotation\": rotations}),\n",
    "        on=\"file_name\",\n",
    "        how=\"left\",\n",
    "    )\n",
    "\n",
    "\n",
    "def update_image():\n",
    "    global updating\n",
    "    updating = True\n",
    "    try:\n",
    "        row = (\n",
    "            get_plant_data(df=table.value, plant_name=sl_plant.value)\n",
    "            .set_index(\"job_id\")\n",
    "            .loc[discrete_player.value]\n",
    "        )\n",
    "        if pt_images.joinpath(row.file_name).is_file() is True:\n",
    "            im_current.object = ti.to_pil(\n",
    "                ti.rotate_image(\n",
    "                    image=ti.load_image(file_name=row.file_name, file_path=pt_images),\n",
    "                    angle=row.rotation,\n",
    "                )\n",
    "            )\n",
    "            ii_disease_index.value = row.di\n",
    "        else:\n",
    "            im_current.object = None\n",
    "    finally:\n",
    "        updating = False\n",
    "\n",
    "\n",
    "@pn.depends(sl_plant.param.value, watch=True)\n",
    "def on_plant_changed(plant_name):\n",
    "    discrete_player.options = get_plant_data(\n",
    "        df=table.value, plant_name=plant_name\n",
    "    ).job_id.to_list()\n",
    "    discrete_player.value = discrete_player.options[0]\n",
    "    update_image()\n",
    "\n",
    "\n",
    "@pn.depends(discrete_player.param.value, watch=True)\n",
    "def on_index_changed(index):\n",
    "    update_image()\n",
    "\n",
    "\n",
    "@pn.depends(ii_disease_index.param.value, watch=True)\n",
    "def on_di_changed(di):\n",
    "    if updating is True:\n",
    "        return\n",
    "    table.value.loc[\n",
    "        (table.value.plant == sl_plant.value)\n",
    "        & (table.value.job_id >= discrete_player.value),\n",
    "        \"di\",\n",
    "    ] = di\n",
    "    table.value.loc[(table.value.plant == sl_plant.value), \"done\"] = True\n",
    "\n",
    "    s_buf = StringIO()\n",
    "    table.value.to_csv(s_buf, sep=\";\")\n",
    "    s_buf.seek(0)\n",
    "    dwn_annotations.filename = f\"{EXPERIMENT}_{dt.now().strftime('%Y%d%m_%H%M%S')}.csv\"\n",
    "    dwn_annotations.label = \"Download\"\n",
    "    dwn_annotations.file = s_buf\n",
    "    dwn_annotations.icon = \"file-download\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build Components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sidebar = pn.Column(\n",
    "    pn.Card(pn.Column(dwn_template, file_input), title=\"File Manager\"),\n",
    "    sl_plant,\n",
    "    pn.Row(\n",
    "        ck_show_finished,\n",
    "        pn.Column(pn.layout.VSpacer(), mk_show_finished, pn.layout.VSpacer()),\n",
    "        height=50,\n",
    "    ),\n",
    ")\n",
    "main = pn.Column(\n",
    "    im_current,\n",
    "    pn.Row(\n",
    "        # pn.layout.HSpacer(),\n",
    "        discrete_player,\n",
    "        ii_disease_index,\n",
    "        dwn_annotations,\n",
    "        # pn.layout.HSpacer(),\n",
    "    ),\n",
    "    max_width=800,\n",
    "    max_height=800,\n",
    "    sizing_mode=\"stretch_width\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pn.Row(sidebar, main)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Render"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template.sidebar.append(sidebar)\n",
    "template.main.append(main)\n",
    "template.servable()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}