{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ralstonia Annotation Tool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load_ext autoreload\n",
    "# %autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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 plotly.express as px\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\", \"plotly\")\n",
    "\n",
    "template = pn.template.BootstrapTemplate(title=\"Ralstonia Annotation Tool\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "EXPERIMENT = \"72AC_PhD_2404\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "INDEX_KEY = \"date_time\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "treatments = {\n",
    "    \"Ralstonia\": \"RS\",\n",
    "    \"Control\": \"CT\",\n",
    "    \"Hydric Stress\": \"HS\",\n",
    "    \"DC3000\": \"DC\",\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## User Interface"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Widgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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",
    "    sizing_mode=\"stretch_width\",\n",
    ")\n",
    "\n",
    "file_input = pn.widgets.FileInput(accept=\".csv,.json\", sizing_mode=\"stretch_width\")\n",
    "table = pn.widgets.Tabulator(\n",
    "    value=pd.DataFrame(),\n",
    "    pagination=\"local\",\n",
    "    page_size=20,\n",
    "    sizing_mode=\"stretch_width\",\n",
    ")\n",
    "\n",
    "\n",
    "sl_treatment = pn.widgets.Select(\n",
    "    name=\"Treatment\", options=list(treatments.keys()), value=\"DC3000\"\n",
    ")\n",
    "sl_plant = pn.widgets.Select(name=\"Plant\", options=[])\n",
    "ck_show_finished = pn.widgets.Checkbox(name=\"Show completed plants\", value=True)\n",
    "ti_done = pn.indicators.BooleanStatus(value=False, color=\"success\")\n",
    "\n",
    "im_current = pn.pane.Plotly(\n",
    "    max_width=800,\n",
    "    max_height=800,\n",
    "    align=(\"center\", \"center\"),\n",
    "    sizing_mode=\"stretch_width\",\n",
    ")\n",
    "discrete_slider = pn.widgets.DiscreteSlider(\n",
    "    name=\"Discrete Player\", options=[0], value=0\n",
    ")\n",
    "bt_previous = pn.widgets.ButtonIcon(icon=\"caret-left\", size=\"4em\", toggle_duration=500)\n",
    "bt_next = pn.widgets.ButtonIcon(icon=\"caret-right\", size=\"4em\", toggle_duration=500)\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=\"Download\",\n",
    "    name=\"Download Annotations\",\n",
    "    sizing_mode=\"stretch_width\",\n",
    "    icon=\"file-download\",\n",
    "    button_type=\"primary\",\n",
    ")\n",
    "pg_completion = pn.indicators.Progress(name=\"Annotation progress\", value=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "updating: bool = False\n",
    "\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",
    "        date_time=lambda x: x.date_time.str.split(\".\", expand=True)[0],\n",
    "    )\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\n",
    "\n",
    "\n",
    "@pn.depends(\n",
    "    table.param.value,\n",
    "    ck_show_finished.param.value,\n",
    "    sl_treatment.param.value,\n",
    "    watch=True,\n",
    ")\n",
    "def on_table_changed(file: str, show_done: bool, treatment):\n",
    "    df = table.value\n",
    "    if \"done\" in df and \"treatment\" in df:\n",
    "        df = df[df.treatment == treatments[treatment]]\n",
    "        sl_plant.options = list(\n",
    "            df.plant.unique()\n",
    "            if show_done is True\n",
    "            else df[df.done == False].plant.unique()\n",
    "        )\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(INDEX_KEY)\n",
    "            .loc[discrete_slider.value]\n",
    "        )\n",
    "        if pt_images.joinpath(row.file_name).is_file() is True:\n",
    "            fig = px.imshow(\n",
    "                ti.rotate_image(\n",
    "                    image=ti.load_image(file_name=row.file_name, file_path=pt_images),\n",
    "                    angle=row.rotation,\n",
    "                ),\n",
    "                width=800,\n",
    "                height=800,\n",
    "            )\n",
    "            fig.update_layout(coloraxis_showscale=False)\n",
    "            fig.update_xaxes(showticklabels=False)\n",
    "            fig.update_yaxes(showticklabels=False)\n",
    "            fig.update_layout(margin=dict(l=0, r=0, t=0, b=0))\n",
    "            im_current.object = fig\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",
    "    plant_data = get_plant_data(df=table.value, plant_name=plant_name)\n",
    "    discrete_slider.options = plant_data[INDEX_KEY].to_list()\n",
    "    ti_done.value = True if plant_data.done.sum() > 0 else False\n",
    "    discrete_slider.value = discrete_slider.options[0]\n",
    "    update_image()\n",
    "\n",
    "\n",
    "@pn.depends(discrete_slider.param.value, watch=True)\n",
    "def on_index_changed(index):\n",
    "    update_image()\n",
    "\n",
    "\n",
    "def do_previous(event):\n",
    "    discrete_slider.value = discrete_slider.options[\n",
    "        max(discrete_slider.options.index(discrete_slider.value) - 1, 0)\n",
    "    ]\n",
    "\n",
    "\n",
    "def do_next(event):\n",
    "    discrete_slider.value = discrete_slider.options[\n",
    "        min(\n",
    "            discrete_slider.options.index(discrete_slider.value) + 1,\n",
    "            len(discrete_slider.options) - 1,\n",
    "        )\n",
    "    ]\n",
    "\n",
    "\n",
    "bt_previous.on_click(do_previous)\n",
    "bt_next.on_click(do_next)\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[INDEX_KEY] >= discrete_slider.value),\n",
    "        \"di\",\n",
    "    ] = di\n",
    "    table.value.loc[(table.value.plant == sl_plant.value), \"done\"] = True\n",
    "    ti_done.value = 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.file = s_buf\n",
    "    dwn_annotations.visible = True\n",
    "\n",
    "    pg_completion.max = len(table.value.plant.unique())\n",
    "    pg_completion.value = len(table.value[table.value.done == True].plant.unique())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build Components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "sidebar = pn.Column(\n",
    "    pn.Card(\n",
    "        pn.Column(\n",
    "            pn.pane.Markdown(\"**Step 1:** Doawload the annaotation template file.\"),\n",
    "            pn.pane.Alert(\n",
    "                \"\"\"Only download the template the first time, changes are stored on your computer only.\"\"\",\n",
    "                alert_type=\"danger\",\n",
    "            ),\n",
    "            dwn_template,\n",
    "            pn.pane.Markdown(\"**Step 2:** Uplaod the downloaded template file.\"),\n",
    "            file_input,\n",
    "            pn.pane.Markdown(\n",
    "                \"\"\"**Step 3:** Annotate images. \n",
    "                Uppon finishing or before closing the app us the download button next to the disease index selector to download the annotations.\"\"\"\n",
    "            ),\n",
    "            pn.pane.Alert(\n",
    "                \"If you want to resume an existing annotation session, upload your latest download.\",\n",
    "                alert_type=\"info\",\n",
    "            ),\n",
    "        ),\n",
    "        title=\"File Manager\",\n",
    "    ),\n",
    "    pn.WidgetBox(\n",
    "        \"#### Plant selection\",\n",
    "        sl_treatment,\n",
    "        sl_plant,\n",
    "        ck_show_finished,\n",
    "        sizing_mode=\"stretch_width\",\n",
    "    ),\n",
    "    pn.WidgetBox(\n",
    "        \"#### Annotation progress\", pg_completion, sizing_mode=\"stretch_width\"\n",
    "    ),\n",
    ")\n",
    "main = pn.Column(\n",
    "    im_current,\n",
    "    pn.Row(bt_previous, discrete_slider, bt_next, ii_disease_index, ti_done, dwn_annotations),\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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}