{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 202311 Dataset Annotation"
   ]
  },
  {
   "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",
    "import warnings\n",
    "from datetime import datetime as dt\n",
    "import inspect\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "import altair as alt\n",
    "import plotly.express as px\n",
    "\n",
    "from PIL import Image, ImageEnhance\n",
    "\n",
    "from siuba import _ as s\n",
    "from siuba import filter as sfilter\n",
    "from siuba import mutate\n",
    "\n",
    "import panel as pn\n",
    "\n",
    "import com_const as cc\n",
    "import com_func as cf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "warnings.simplefilter(action=\"ignore\", category=UserWarning)\n",
    "warnings.simplefilter(action=\"ignore\", category=FutureWarning)"
   ]
  },
  {
   "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\", 16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alt.data_transformers.disable_max_rows();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pn.extension(\n",
    "    \"plotly\", \"terminal\", \"tabulator\", \"vega\", notifications=True, console_output=\"disable\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template = pn.template.BootstrapTemplate(title=\"OIV Annotation Tool\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gobals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_row = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_quality_options = {\n",
    "    \"?\":np.nan,\n",
    "    \"good\":\"good_images\",\n",
    "    \"crop\":\"crop_images\",\n",
    "    \"missing\":\"missing_images\",\n",
    "    \"dark\":\"dark_images\",\n",
    "    \"blur\":\"blur_images\",\n",
    "    \"color\":\"color_images\",\n",
    "    \"water\":\"water_images\",\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = cf.read_dataframe(path=cc.path_to_data.joinpath(\"oiv_annotation.csv\")).sort_values(\n",
    "    [\"experiment\", \"inoc\", \"dpi\", \"plaque\", \"row\", \"col\"]\n",
    ")\n",
    "if \"seen_at\" not in df:\n",
    "    df = df >> mutate(seen_at=np.nan)\n",
    "df.seen_at = pd.to_datetime(df.seen_at)\n",
    "df = df.set_index(\"file_name\")\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_image(image_name:str, color, brightness, contrast, sharpness):\n",
    "    image_path = cc.path_to_leaf_patches.joinpath(image_name)\n",
    "    if image_path.is_file() is False:\n",
    "        fig = px.imshow(\n",
    "            np.array(\n",
    "                [\n",
    "                    [[255, 0, 255], [255, 0, 255], [255, 0, 255]],\n",
    "                    [[255, 0, 255], [255, 0, 255], [255, 0, 255]],\n",
    "                    [[255, 0, 255], [255, 0, 255], [255, 0, 255]],\n",
    "                ],\n",
    "                dtype=np.uint8,\n",
    "            )\n",
    "        )\n",
    "    else:\n",
    "        image = Image.open(image_path)\n",
    "        image = ImageEnhance.Color(image=image).enhance(color)\n",
    "        image = ImageEnhance.Contrast(image=image).enhance(contrast)\n",
    "        image = ImageEnhance.Brightness(image=image).enhance(brightness)\n",
    "        image = ImageEnhance.Sharpness(image=image).enhance(sharpness)\n",
    "        fig = px.imshow(image)\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",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_classes(df_: pd.DataFrame, var: str):\n",
    "    d = pd.DataFrame(\n",
    "        data={\n",
    "            var: df_[var]\n",
    "            .fillna(\"?\")\n",
    "            .astype(str)\n",
    "            .str.replace(\".0\", \"\")\n",
    "            .str.replace(\"images\", \"\")\n",
    "        }\n",
    "    )\n",
    "    bars = (\n",
    "        alt.Chart(d)\n",
    "        .mark_bar()\n",
    "        .encode(\n",
    "            y=alt.Y(var, title=None),\n",
    "            x=alt.X(\"count()\", axis=None),\n",
    "            color=alt.Color(var, legend=None),\n",
    "        )\n",
    "    )\n",
    "    text = bars.mark_text(align=\"center\", dy=0, dx=12).encode(\n",
    "        y=alt.Y(var, title=None),\n",
    "        x=alt.X(\"count()\", axis=None),\n",
    "        color=alt.Color(var, legend=None),\n",
    "        text=\"count()\",\n",
    "    )\n",
    "\n",
    "    return (bars + text).configure_view(stroke=None).configure_axis(grid=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Widgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_current = pn.pane.Plotly(height=750, align=(\"center\", \"center\"))\n",
    "mkd_current = pn.pane.Markdown(sizing_mode=\"scale_width\", align=\"center\")\n",
    "\n",
    "sl_contrast = pn.widgets.EditableFloatSlider(\n",
    "    name=\"Contrast\", start=0.0, end=7.5, value=1.5, step=0.1, sizing_mode=\"scale_width\"\n",
    ")\n",
    "sl_color = pn.widgets.EditableFloatSlider(\n",
    "    name=\"Color\", start=0.0, end=5.0, value=1.0, step=0.1, sizing_mode=\"scale_width\"\n",
    ")\n",
    "sl_brightness = pn.widgets.EditableFloatSlider(\n",
    "    name=\"Brightness\",\n",
    "    start=0.0,\n",
    "    end=5.0,\n",
    "    value=1.0,\n",
    "    step=0.1,\n",
    "    sizing_mode=\"scale_width\",\n",
    ")\n",
    "sl_sharpness = pn.widgets.EditableFloatSlider(\n",
    "    name=\"Sharpness\", start=0.0, end=2.0, value=1.5, step=0.1, sizing_mode=\"scale_width\"\n",
    ")\n",
    "\n",
    "c_image_processing = pn.Card(\n",
    "    pn.Column(sl_brightness, sl_color, sl_contrast, sl_sharpness),\n",
    "    title=\"Image Processing Options\",\n",
    "    sizing_mode=\"scale_width\",\n",
    ")\n",
    "\n",
    "pg_progress = pn.widgets.Tqdm(name=\"Progress\", align=\"center\", max=len(df))\n",
    "\n",
    "rgb_oiv = pn.widgets.RadioButtonGroup(\n",
    "    name=\"OIV\",\n",
    "    options=[\"?\", 1, 3, 5, 7, 9],\n",
    "    button_style=\"outline\",\n",
    "    button_type=\"success\",\n",
    ")\n",
    "\n",
    "rgb_source = pn.widgets.RadioButtonGroup(\n",
    "    name=\"Image quality\",\n",
    "    options=list(image_quality_options.keys()),\n",
    "    button_style=\"outline\",\n",
    "    button_type=\"success\",\n",
    "    value=\"?\",\n",
    ")\n",
    "\n",
    "sel_def_img_quality = pn.widgets.Select(\n",
    "    name=\"Default Image Quality\", options=list(image_quality_options.keys())\n",
    ")\n",
    "\n",
    "mc_filter_quality = pn.widgets.MultiChoice(\n",
    "    name=\"Allow qualities\",\n",
    "    options=list(image_quality_options.values()),\n",
    "    value=list(image_quality_options.values()),\n",
    ")\n",
    "\n",
    "rgb_target = pn.widgets.RadioButtonGroup(\n",
    "    name=\"Annotation target\",\n",
    "    options=[\"All\", \"OIV\", \"Image quality\"],\n",
    "    button_style=\"outline\",\n",
    "    button_type=\"success\",\n",
    "    value=\"All\",\n",
    ")\n",
    "\n",
    "c_anno_options = pn.Card(\n",
    "    pn.Column(\n",
    "        pn.Row(pn.pane.Markdown(\"**Annotate**\"), rgb_target),\n",
    "        sel_def_img_quality,\n",
    "        mc_filter_quality,\n",
    "    ),\n",
    "    title=\"Annotation Options\",\n",
    "    sizing_mode=\"scale_width\",\n",
    ")\n",
    "\n",
    "pn_hist_oiv = pn.pane.Vega()\n",
    "pn_hist_source = pn.pane.Vega()\n",
    "\n",
    "c_hists = pn.Card(\n",
    "    pn.Column(\n",
    "        pn.pane.Markdown(\"### OIV\"),\n",
    "        pn_hist_oiv,\n",
    "        pn.pane.Markdown(\"### Image Quality\"),\n",
    "        pn_hist_source,\n",
    "    ),\n",
    "    title=\"Annotation Overview\",\n",
    "    sizing_mode=\"scale_width\",\n",
    ")\n",
    "\n",
    "sw_ui_state = pn.widgets.Switch(name=\"active\", value=False)\n",
    "alt_ui_state = pn.pane.Alert(\"Annotations will be stored\", alert_type=\"primary\")\n",
    "\n",
    "pn_ui_state = pn.Row(sw_ui_state, alt_ui_state)\n",
    "\n",
    "\n",
    "bt_next = pn.widgets.Button(name=\"Next\", button_type=\"primary\")\n",
    "bt_previous = pn.widgets.Button(name=\"Previous\", button_type=\"primary\")\n",
    "\n",
    "ui_annotation = pn.GridSpec(sizing_mode=\"scale_width\", align=\"center\", max_height=120)\n",
    "\n",
    "ui_annotation[1, 0] = bt_previous\n",
    "ui_annotation[0, 1:5] = rgb_source\n",
    "ui_annotation[1, 1:5] = rgb_oiv\n",
    "ui_annotation[1, 5] = bt_next"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@pn.depends(\n",
    "    sl_color.param.value,\n",
    "    sl_contrast.param.value,\n",
    "    sl_brightness.param.value,\n",
    "    sl_sharpness.param.value,\n",
    "    watch=True,\n",
    ")\n",
    "def on_preprocess_changed(color, contrast, brightness, sharpeness):\n",
    "    img_current.object = update_image(\n",
    "        image_name=current_row.file_name,\n",
    "        color=color,\n",
    "        brightness=brightness,\n",
    "        contrast=contrast,\n",
    "        sharpness=sharpeness,\n",
    "    )\n",
    "\n",
    "\n",
    "def update_ui_state(ui_state: bool):\n",
    "    if ui_state is True:\n",
    "        alt_ui_state.object = \"Annotations will be stored\"\n",
    "        alt_ui_state.alert_type = \"primary\"\n",
    "    else:\n",
    "        alt_ui_state.object = \"Annotations will be discarded\"\n",
    "        alt_ui_state.alert_type = \"danger\"\n",
    "\n",
    "\n",
    "@pn.depends(sw_ui_state, watch=True)\n",
    "def on_ui_State_changed(new_state: bool):\n",
    "    update_ui_state(new_state)\n",
    "\n",
    "\n",
    "def select_next(event):\n",
    "    global current_row\n",
    "    global df\n",
    "    now = dt.now()\n",
    "    if current_row is not None and (event is None or event.obj.name == \"Next\"):\n",
    "        if rgb_target.value in [\"All\", \"OIV\"] and rgb_oiv.value != \"?\":\n",
    "            df.at[current_row.file_name, \"oiv\"] = int(rgb_oiv.value)\n",
    "            df.at[current_row.file_name, \"oiv_annotated_at\"] = now\n",
    "\n",
    "        if rgb_target.value in [\"All\", \"Image quality\"] and rgb_source.value != \"?\":\n",
    "            df.at[current_row.file_name, \"source_annotated_at\"] = now\n",
    "            df.at[current_row.file_name, \"source\"] = image_quality_options[\n",
    "                rgb_source.value\n",
    "            ]\n",
    "        cf.write_dataframe(\n",
    "            df=df.reset_index(),\n",
    "            path=cc.path_to_data.joinpath(\n",
    "                \"oiv_annotation.csv\" if sw_ui_state.value is True else \"oiv_annotation_test.csv\"\n",
    "            ),\n",
    "        )\n",
    "        df.at[current_row.file_name, \"seen_at\"] = now\n",
    "\n",
    "    df_cr = df >> sfilter(s.source.isin(mc_filter_quality.value))\n",
    "\n",
    "    if rgb_target.value == \"All\":\n",
    "        df_cr = df_cr >> sfilter(s.oiv.isna() | s.source.isna())\n",
    "    elif rgb_target.value == \"OIV\":\n",
    "        df_cr = df_cr >> sfilter(s.oiv.isna())\n",
    "    if rgb_target.value == \"Image quality\":\n",
    "        df_cr = df_cr >> sfilter(s.source.isna())\n",
    "    remaining = len(df_cr)\n",
    "    if event is None or event.obj.name == \"Next\":\n",
    "        df_cr = df_cr.reset_index()\n",
    "        current_row = df_cr.sample(n=1).iloc[0] if len(df_cr) > 0 else None\n",
    "    elif event.obj.name == \"Previous\":\n",
    "        current_row = (\n",
    "            (df.reset_index() >> sfilter(~s.seen_at.isna()))\n",
    "            .sort_values(\"seen_at\", ascending=False)\n",
    "            .reset_index(drop=True)\n",
    "            .iloc[0]\n",
    "        )\n",
    "        df.at[current_row.file_name, \"seen_at\"] = None\n",
    "\n",
    "    if current_row is not None:\n",
    "        rgb_source.value = (\n",
    "            sel_def_img_quality.value\n",
    "            if pd.isnull(current_row.source)\n",
    "            else {v: k for k, v in image_quality_options.items()}[current_row.source]\n",
    "        )\n",
    "        rgb_oiv.value = (\n",
    "            current_row.oiv if current_row.oiv in [1, 3, 5, 7, 9] else \"?\"\n",
    "        )\n",
    "\n",
    "    pg_progress.value = len(df) - remaining\n",
    "    file_name = current_row.file_name if current_row is not None else \"\"\n",
    "    img_current.object = update_image(\n",
    "        image_name=file_name,\n",
    "        color=sl_color.value,\n",
    "        brightness=sl_brightness.value,\n",
    "        contrast=sl_contrast.value,\n",
    "        sharpness=sl_sharpness.value,\n",
    "    )\n",
    "    mkd_current.object = f\"## {file_name}\"\n",
    "    df_unf = df >> sfilter(s.source.isin(mc_filter_quality.value))\n",
    "    pn_hist_source.object = plot_classes(df_unf, \"source\")\n",
    "    pn_hist_oiv.object = plot_classes(df_unf, \"oiv\")\n",
    "\n",
    "\n",
    "@pn.depends(rgb_target, watch=True)\n",
    "def on_target_changed(target):\n",
    "    rgb_oiv.disabled = target == \"Image quality\"\n",
    "    rgb_source.disabled = target == \"OIV\"\n",
    "\n",
    "\n",
    "# @pn.depends(rgb_oiv, watch=True)\n",
    "# def on_oiv_changed(_):\n",
    "#     select_next(None)\n",
    "\n",
    "\n",
    "bt_next.on_click(select_next)\n",
    "bt_previous.on_click(select_next)\n",
    "\n",
    "update_ui_state(sw_ui_state.value)\n",
    "select_next(None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## UI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template.sidebar.append(pn_ui_state)\n",
    "template.sidebar.append(c_image_processing)\n",
    "template.sidebar.append(c_anno_options)\n",
    "\n",
    "template.main.append(\n",
    "    pn.Row(\n",
    "        pn.Column(\n",
    "            # mkd_current,\n",
    "            img_current,\n",
    "            ui_annotation,\n",
    "        ),\n",
    "        pn.Column(c_hists, pg_progress),\n",
    "    )\n",
    ")\n",
    "\n",
    "template.servable()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Please launch with command \"panel serve leaf_patch_annotation.ipynb --show --dev\" from the \"src\" folder"
   ]
  }
 ],
 "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}