webview/webview.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# WebViewer\n",
    "\n",
    "This notebook exemplifies the visualization with the webviewer tool. It can also be served through port forwarding so that the data can be visualized remotely. For more information on how to dop this please check the Readme.\n",
    "\n",
    "> **important**: Please run the Example datasets first before running this notebook, as it expects the tree and mesh outputs to exist! \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import panel as pn\n",
    "pn.extension('vtk')\n",
    "pn.extension()\n",
    "import holoviews as hv\n",
    "hv.extension('bokeh', width=100)\n",
    "\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import param\n",
    "import matplotlib.pyplot as plt\n",
    "from holoviews import opts\n",
    "from holoviews.streams import Selection1D\n",
    "from inter_view.color import clipped_plasma_r, plot_colorbar\n",
    "from improc.io import parse_collection, DCAccessor\n",
    "from lstree.lineage.plot import plot_tree\n",
    "import time\n",
    "from io import BytesIO\n",
    "import argparse\n",
    "\n",
    "DCAccessor.register()\n",
    "\n",
    "\n",
    "opts.defaults(\n",
    "    opts.Layout(sizing_mode='stretch_both'),\n",
    "    opts.Overlay('tree', yaxis=None, show_title=False, framewise=True, responsive=True, min_width=1000, min_height=1000),\n",
    "    opts.Points('tree', size=3, cmap=clipped_plasma_r, framewise=True, colorbar=True, colorbar_position='right', responsive=True, colorbar_opts={'bar_line_width':0, 'width':15}),\n",
    "    opts.Segments(line_width=1, color='black', framewise=True, responsive=True),\n",
    "    opts.VLine(line_dash='dashed', color='grey', line_width=1, responsive=True, framewise=True))\n",
    "\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\"--basedir\", help=\"data directory\", nargs=1)\n",
    "args, unknown = parser.parse_known_args()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parse meshes and trees\n",
    "\n",
    "First we parse all of the mesh and tree outputs we got from the two example datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# when using this viewer as a notebook, specify base data directory here. Currently the folder with the Example datasets is set. When using remotely, the basedir argument must be provided.\n",
    "if args.basedir is None:\n",
    "    basedir = '../example/data'\n",
    "else:\n",
    "    basedir = args.basedir[0]\n",
    "\n",
    "\n",
    "# create/update symlink to access mesh data from javascript\n",
    "static_basedir = 'static_data'\n",
    "if os.path.exists(static_basedir):\n",
    "    os.unlink(static_basedir)\n",
    "    \n",
    "os.symlink(basedir, static_basedir, target_is_directory=True)\n",
    "\n",
    "df_grid = parse_collection(os.path.join(static_basedir, '{dataset}/{subdir}/{f2}T{time:04d}.vti'), ['dataset', 'subdir', 'time'])\n",
    "df_mesh = parse_collection(os.path.join(static_basedir, '{dataset}/{subdir}/{f2}T{time:04d}.vtp'), ['dataset', 'subdir', 'time'])\n",
    "df_vtk = pd.concat([df_grid, df_mesh])\n",
    "df_h5 = parse_collection(static_basedir + '/{dataset}/agg_features.{ext}', ['dataset'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load first sample and get list of available node/edge attributes\n",
    "This is just to check that indeed everything fits. We load the first dataset from the list, and show all of its node properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfn = pd.read_hdf(df_h5.iloc[0:1].dc.path[0], 'nodes')\n",
    "dfe = pd.read_hdf(df_h5.iloc[0:1].dc.path[0], 'edges')\n",
    "\n",
    "NODE_ATTRIBUTES = sorted(dfn.select_dtypes(['int', 'float']).columns.tolist() + ['selection', 'neighbor_trace'])\n",
    "EDGE_ATTRIBUTES = sorted(list( set(dfe.select_dtypes(['int', 'float']).columns) - {'node_start', 'node_stop'}))\n",
    "\n",
    "dfn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tree viewer\n",
    "This class is used for displaying the trees that we have in our datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class TreeViewer(param.Parameterized):\n",
    "    '''Tree viewer for interactive visualization with Holoviews'''\n",
    "    time_interval = param.Number(10/60)\n",
    "    h5_path = param.Parameter()\n",
    "    dfe = param.DataFrame()\n",
    "    dfn = param.DataFrame()\n",
    "    \n",
    "    node_color = param.ObjectSelector(default='dodgerblue', objects=['dodgerblue'] + NODE_ATTRIBUTES)\n",
    "    node_size = param.ObjectSelector(default='fixed', objects=['fixed', None] + NODE_ATTRIBUTES)\n",
    "    \n",
    "    edge_color = param.ObjectSelector(default='black', objects=['black'] + EDGE_ATTRIBUTES)\n",
    "    edge_size = param.ObjectSelector(default='fixed', objects=['fixed'] + EDGE_ATTRIBUTES)\n",
    "    \n",
    "    tooltips = param.List(['mamut_id', 'mamut_t'])\n",
    "    \n",
    "    hv_tree = param.Parameter(hv.Overlay([hv.Scatter([])]))\n",
    "    \n",
    "    last_tap = param.Parameter(hv.streams.SingleTap(transient=True), instantiate=True)\n",
    "    last_clicked_mamut_t = param.Number(1)\n",
    "    tree_update_counter = param.Integer(0)\n",
    "    \n",
    "    selection = param.Parameter(Selection1D(), instantiate=True)\n",
    "    selection_update_counter = param.Integer(0)\n",
    "    \n",
    "    def __init__(self, *args, **kargs):\n",
    "        super().__init__(*args, **kargs)\n",
    "        \n",
    "        self._load_tree()\n",
    "    \n",
    "    def _set_time_interval(self):\n",
    "        '''estimate time interval upon loading new data'''\n",
    "        self.time_interval = np.median(np.diff(np.sort(self.dfn.time.unique())))\n",
    "    \n",
    "    @param.depends('h5_path', watch=True)\n",
    "    def _load_tree(self):\n",
    "        if self.h5_path:\n",
    "            \n",
    "            try:\n",
    "                self.dfn = pd.read_hdf(self.h5_path, 'nodes')#.iloc[:100]\n",
    "                self.dfn['selection']  = 0\n",
    "                self.dfn['neighbor_trace']  = 0\n",
    "                self.dfe = pd.read_hdf(self.h5_path, 'edges')#.iloc[:0]\n",
    "                self._set_time_interval()\n",
    "            except KeyError as e:\n",
    "                self.dfn = None\n",
    "                \n",
    "            self.tree_update_counter += 1\n",
    "    \n",
    "    @param.depends('selection.index', watch=True)\n",
    "    def update_selection_column(self):\n",
    "        if len(self.selection.index) > 0:\n",
    "            self.dfn['selection']  = 0\n",
    "            self.dfn['selection'].iloc[self.selection.index] = 1\n",
    "            self.selection_update_counter += 1\n",
    "            \n",
    "            if 'selection' in [self.node_color, self.node_size, self.edge_color, self.edge_size]:\n",
    "                self.tree_update_counter += 1\n",
    "    \n",
    "    @param.depends('tree_update_counter', 'node_color', 'node_size', 'edge_color', 'edge_size')\n",
    "    def _plot_tree(self):\n",
    "        \n",
    "        try:\n",
    "            # dynamic update of color,size,etc. does not work when it depends on vdims --> we are forced to rebuild \n",
    "            # the plot everytime...\n",
    "            if self.node_size == 'fixed':\n",
    "                node_size = 5\n",
    "            else:\n",
    "                node_size = self.node_size\n",
    "\n",
    "            if self.edge_size == 'fixed':\n",
    "                edge_size = 2\n",
    "            else:\n",
    "                edge_size = self.edge_size\n",
    "\n",
    "            self.hv_tree = plot_tree(self.dfn,\n",
    "                                     self.dfe,\n",
    "                                     tooltips=self.tooltips,\n",
    "                                     node_size=node_size,\n",
    "                                     node_color=self.node_color,\n",
    "                                     edge_width=edge_size,\n",
    "                                     edge_color=self.edge_color,\n",
    "                                     min_node_size=2,\n",
    "                                     max_node_size=30,\n",
    "                                     min_edge_width=1,\n",
    "                                     max_edge_width=9,)\n",
    "\n",
    "            # link the last clicked position\n",
    "            self.last_tap.source = self.hv_tree.Tree.II\n",
    "\n",
    "            # add selection stream and attach callback to update sample/image selection\n",
    "            self.selection.source=self.hv_tree.Tree.II\n",
    "\n",
    "            return self.hv_tree\n",
    "        \n",
    "        except Exception as e:\n",
    "            # tree does not exist or does not contain the requried features\n",
    "            self.hv_tree = hv.Overlay([hv.Scatter([])])\n",
    "            self.dfn = None\n",
    "            self.dfe = None\n",
    "            return self.hv_tree\n",
    "    \n",
    "    # somehow does not work when combined with vtk pane\n",
    "    @param.depends('last_tap.x', watch=True)\n",
    "    def _update_last_clicked_timepoint(self):\n",
    "        if self.last_tap.x is not None:\n",
    "            self.last_clicked_mamut_t = int(round(self.last_tap.x / self.time_interval))\n",
    "    \n",
    "    def widgets(self):\n",
    "        return pn.WidgetBox(self.param.node_color,\n",
    "                            self.param.node_size,\n",
    "                            self.param.edge_color,\n",
    "                            self.param.edge_size)\n",
    "    \n",
    "    def plot(self):    \n",
    "        tree_dmap = hv.DynamicMap(self._plot_tree)#.opts(responsive=True)\n",
    "        return tree_dmap.opts(xlabel='time').opts(opts.Points(tools=['hover', 'box_select', 'lasso_select', 'tap']))\n",
    "    \n",
    "    def panel(self):\n",
    "        # dummy plot somehow needed forthe selection to continue updating after dynamic map is refreshed\n",
    "        self.dummy = hv.Points([], group='tree')\n",
    "        return pn.Row(self.widgets(), self.dummy * self.plot())\n",
    "\n",
    "\n",
    "# ds_id = 6\n",
    "# tv = TreeViewer(h5_path=df_h5.dc[ds_id].dc.path[0],\n",
    "#                 node_color='nuclei_volume')\n",
    "\n",
    "# tv.panel().servable()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HTML Panel\n",
    "Needed for sending arbitrary javascript code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class JSCode():\n",
    "    '''Creates a (hidden) dummy HTML panel that allows \n",
    "    sending arbitrary javascript code'''\n",
    "    \n",
    "    def __init__(self):\n",
    "        self.dummy = pn.pane.HTML(\"\", style={'padding':'0px', 'visibility':'hidden'}, width=0, height=0, margin=0)\n",
    "    \n",
    "    def run(self, code):\n",
    "        self.dummy.object = ''\n",
    "        self.dummy.object = \"\"\"<script type=\"text/javascript\">{}</script>\"\"\".format(code)\n",
    "\n",
    "def make_feature_lut(cmap, ids, values, bounds):\n",
    "    ids = np.array(ids)[:,None]\n",
    "    values = (values-bounds[0]) / (bounds[1]-bounds[0])\n",
    "    colors = cmap(values)[:,:3]\n",
    "    return np.concatenate([ids, colors], axis=1).tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Organoid viewer\n",
    "Similar to the Tree viewer, but this time for visualization of the vtk mesh outputs form the segmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "class OrganoidViewer(param.Parameterized):\n",
    "    '''Visualization of 3D segmetnation meshes.'''\n",
    "    \n",
    "    spacing = param.Tuple((1,1,1))\n",
    "    \n",
    "    nuclei_mesh = param.Parameter()\n",
    "    cell_mesh = param.Parameter()\n",
    "    rgb_grid = param.Parameter()\n",
    "    \n",
    "    features = param.DataFrame()\n",
    "    features_bounds = param.Dict()\n",
    "\n",
    "    cell_representation = param.ObjectSelector(default=\"Wireframe\", objects=['Wireframe', 'Surface'])\n",
    "    nuclei_representation = param.ObjectSelector(default=\"Surface\", objects=['Wireframe', 'Surface'])\n",
    "    \n",
    "    cell_opacity = param.Number(0.1, bounds=(0,1), step=0.01)\n",
    "    nuclei_opacity = param.Number(1., bounds=(0,1), step=0.01)\n",
    "    \n",
    "    cell_color = param.ObjectSelector(default=None, objects=[None] + NODE_ATTRIBUTES)\n",
    "    nuclei_color = param.ObjectSelector(default=None, objects=[None] + NODE_ATTRIBUTES)\n",
    "    cell_lut = param.List([[0, 1., 1., 1.]])\n",
    "    nuclei_lut = param.List([[0, 0., 0.56470588, 1.]])\n",
    "    \n",
    "    cmap = param.Parameter(clipped_plasma_r)\n",
    "    \n",
    "    def __init__(self, *args, **kargs):\n",
    "        super().__init__(*args, **kargs)\n",
    "        \n",
    "        self.jscode = JSCode()\n",
    "        \n",
    "        self.nuclei_opacity_wg = pn.panel(self.param.nuclei_opacity)\n",
    "        self.cell_opacity_wg = pn.panel(self.param.cell_opacity)\n",
    "        self.nuclei_representation_wg = pn.panel(self.param.nuclei_representation)\n",
    "        self.cell_representation_wg = pn.panel(self.param.cell_representation)\n",
    "        \n",
    "        with open('vtkjs-panel.html', 'r') as f:\n",
    "            html_vtkjs_panel = f.read()\n",
    "            \n",
    "        # set paths to load on first render\n",
    "        if self.features is not None:\n",
    "            self._update_nuclei_color()\n",
    "            self._update_cell_color()\n",
    "            code = self._generate_update_js_code()\n",
    "        else:\n",
    "            code = ''\n",
    "        html_vtkjs_panel = html_vtkjs_panel.replace('{{init_code}}',code)\n",
    "\n",
    "        self.html_pane = pn.pane.HTML(html_vtkjs_panel, min_height=800, min_width=800, sizing_mode='stretch_both')\n",
    "        self._set_callbacks()\n",
    "\n",
    "        \n",
    "    def clear_renderer(self):\n",
    "        self.jscode.run(\"\"\"renderer.removeActor(imageActorK);\n",
    "                           renderer.removeActor(imageActorI);\n",
    "                           renderer.removeActor(imageActorJ);\n",
    "                           renderer.removeActor(nucleiActor);\n",
    "                           renderer.removeActor(cellActor);\n",
    "                           render();\"\"\")\n",
    "    \n",
    "    def _generate_update_js_code(self):\n",
    "        \n",
    "        code = \"\"\"\n",
    "               nucleiPath = '{nuclei_path}';\n",
    "               var colorArray = JSON.parse('{nuclei_lut}');\n",
    "               updateNucleiMesh(colorArray, {nuclei_opacity}, '{nuclei_repr}');\n",
    "               \n",
    "               cellPath = '{cell_path}';\n",
    "               var colorArray = JSON.parse('{cell_lut}');\n",
    "               updateCellMesh(colorArray, {cell_opacity}, '{cell_repr}');\n",
    "               \n",
    "               imgPath = '{grid_path}'\n",
    "               updateImage();\n",
    "               \"\"\".format(\n",
    "                    nuclei_path=self.nuclei_mesh,\n",
    "                    nuclei_lut=self.nuclei_lut,\n",
    "                    nuclei_opacity=self.nuclei_opacity_wg.value,\n",
    "                    nuclei_repr=self.nuclei_representation_wg.value,\n",
    "                    cell_path=self.cell_mesh,\n",
    "                    cell_lut=self.cell_lut,\n",
    "                    cell_opacity=self.cell_opacity_wg.value,\n",
    "                    cell_repr=self.cell_representation_wg.value,\n",
    "                    grid_path=self.rgb_grid)\n",
    "\n",
    "        \n",
    "        return code\n",
    "        \n",
    "    def load(self, features, nuclei_mesh, cell_mesh, rgb_grid=None, features_bounds=None):\n",
    "        self.features = features\n",
    "        self.nuclei_mesh = nuclei_mesh\n",
    "        self.cell_mesh = cell_mesh\n",
    "        self.rgb_grid = rgb_grid\n",
    "                \n",
    "        if features_bounds is not None:\n",
    "            self.features_bounds = features_bounds\n",
    "        \n",
    "        self._update_nuclei_lut()\n",
    "        self._update_cell_lut()\n",
    "        \n",
    "        code = self._generate_update_js_code()\n",
    "        self.jscode.run(code)\n",
    "        \n",
    "    def update_features(self, features):\n",
    "        if self.nuclei_color in self.features.columns:\n",
    "            nuclei_changed = not self.features[self.nuclei_color].equals(features[self.nuclei_color])\n",
    "        else:\n",
    "            nuclei_changed = False\n",
    "            \n",
    "        if self.cell_color in self.features.columns:\n",
    "            cell_changed = not self.features[self.cell_color].equals(features[self.cell_color])\n",
    "        else:\n",
    "            cell_changed = False\n",
    "        \n",
    "        self.features = features\n",
    "        if nuclei_changed:\n",
    "            self._update_nuclei_color()\n",
    "            \n",
    "        if cell_changed:\n",
    "            self._update_cell_color()\n",
    "        \n",
    "    def _set_callbacks(self):\n",
    "        \n",
    "        self.nuclei_opacity_wg.jslink(self.html_pane , code={'value':\"\"\"\n",
    "            nucleiActor.getProperty().setOpacity(source.value);\n",
    "            render();\n",
    "            \"\"\"})\n",
    "        \n",
    "        self.cell_opacity_wg.jslink(self.html_pane , code={'value':\"\"\"\n",
    "            cellActor.getProperty().setOpacity(source.value);\n",
    "            render();\n",
    "            \"\"\"})\n",
    "        \n",
    "        self.nuclei_representation_wg.jslink(self.html_pane , code={'value':\"\"\"\n",
    "            if(source.value==\"Wireframe\"){{\n",
    "                nucleiActor.getProperty().setRepresentationToWireframe();\n",
    "            }}else if(source.value==\"Surface\"){{\n",
    "                nucleiActor.getProperty().setRepresentationToSurface();\n",
    "            }}\n",
    "            render();\n",
    "            \"\"\"})\n",
    "        \n",
    "        self.cell_representation_wg.jslink(self.html_pane , code={'value':\"\"\"\n",
    "            if(source.value==\"Wireframe\"){{\n",
    "                cellActor.getProperty().setRepresentationToWireframe();\n",
    "            }}else if(source.value==\"Surface\"){{\n",
    "                cellActor.getProperty().setRepresentationToSurface();\n",
    "            }}\n",
    "            render();\n",
    "            \"\"\"})\n",
    "    \n",
    "    def _update_nuclei_lut(self):\n",
    "        if self.nuclei_color:\n",
    "            self.nuclei_lut = make_feature_lut(cmap=self.cmap,\n",
    "                                   ids=self.features['timepoint_id']+1,\n",
    "                                   values=self.features[self.nuclei_color],\n",
    "                                   bounds=self.features_bounds[self.nuclei_color])\n",
    "        else:\n",
    "            self.nuclei_lut = [[0, 0., 0.56470588, 1.]]\n",
    "    \n",
    "    @param.depends('nuclei_color', watch=True)\n",
    "    def _update_nuclei_color(self):\n",
    "        \n",
    "        self._update_nuclei_lut()\n",
    "        \n",
    "        self.jscode.run(\"\"\"\n",
    "            var colorArray = JSON.parse('{}')\n",
    "            updateLUT(nucleiLookupTable, colorArray, nucleiPolydata);\n",
    "        \"\"\".format(str(self.nuclei_lut)))\n",
    "        \n",
    "        \n",
    "    def _update_cell_lut(self):\n",
    "        if self.cell_color:\n",
    "            self.cell_lut = make_feature_lut(cmap=self.cmap,\n",
    "                                   ids=self.features['timepoint_id']+1,\n",
    "                                   values=self.features[self.cell_color],\n",
    "                                   bounds=self.features_bounds[self.cell_color])\n",
    "        else:\n",
    "            self.cell_lut = [[0, 1., 1., 1.]]\n",
    "    \n",
    "    @param.depends('cell_color', watch=True)\n",
    "    def _update_cell_color(self):\n",
    "        \n",
    "        self._update_cell_lut()\n",
    "        \n",
    "        self.jscode.run(\"\"\"\n",
    "            var colorArray = JSON.parse('{}')\n",
    "            updateLUT(cellLookupTable, colorArray, nucleiPolydata);\n",
    "        \"\"\".format(str(self.cell_lut)))\n",
    "    \n",
    "    def _bg_hook(self, plot, element):\n",
    "        '''Directly access bokeh figure and set selection'''\n",
    "        plot.state.border_fill_color = \"#f5f5f5\"\n",
    "    \n",
    "    @param.depends('nuclei_color')\n",
    "    def _nuclei_colorbar(self):\n",
    "        if self.nuclei_color is None:\n",
    "            return None\n",
    "        else:\n",
    "            cbar = plot_colorbar(clipped_plasma_r, self.features_bounds[self.nuclei_color], orientation='horizontal', backend='bokeh')\n",
    "            return cbar.opts(hooks=[self._bg_hook])\n",
    "\n",
    "    @param.depends('cell_color')\n",
    "    def _cell_colorbar(self):\n",
    "        if self.cell_color is None:\n",
    "            return None\n",
    "        else:\n",
    "            cbar = plot_colorbar(clipped_plasma_r, self.features_bounds[self.cell_color], orientation='horizontal', backend='bokeh')\n",
    "            return cbar.opts(hooks=[self._bg_hook])\n",
    "\n",
    "    def _nuclei_widgets(self):\n",
    "        return pn.WidgetBox(self.nuclei_opacity_wg,\n",
    "                            self.param.nuclei_color,\n",
    "                            self._nuclei_colorbar,\n",
    "                            self.nuclei_representation_wg,)\n",
    "        \n",
    "    def _cell_widgets(self):\n",
    "        return pn.WidgetBox(self.cell_opacity_wg,\n",
    "                            self.param.cell_color,\n",
    "                            self._cell_colorbar,\n",
    "                            self.cell_representation_wg,)\n",
    "        \n",
    "    def widgets(self):\n",
    "        return pn.Row(self._nuclei_widgets(),\n",
    "                      self._cell_widgets(),\n",
    "#                       pn.Column(self._slice_widgets, self._colorbars)\n",
    "                     )\n",
    "    \n",
    "    @property\n",
    "    def vtkpan(self):\n",
    "        return pn.Row(self.html_pane, self.jscode.dummy, sizing_mode='stretch_both')\n",
    "    \n",
    "    def panel(self):\n",
    "        return pn.Column(self.vtkpan,\n",
    "                         self.widgets())\n",
    "\n",
    "\n",
    "# mamut_t = 323\n",
    "\n",
    "# ds_id = 3\n",
    "# h5_path = df_h5.dc[ds_id].dc.path[0]\n",
    "# dfn = pd.read_hdf(h5_path, 'nodes')\n",
    "# dfe = pd.read_hdf(h5_path, 'edges')\n",
    "\n",
    "# features_bounds = dfn.agg(['min', 'max']).to_dict('list')\n",
    "\n",
    "# db = OrganoidViewer(features=dfn.set_index('mamut_t').loc[mamut_t:mamut_t],\n",
    "#                     features_bounds=features_bounds,\n",
    "#                     nuclei_mesh=df_vtk.dc[ds_id, 'nuclei_mesh', mamut_t].dc.path[0],\n",
    "#                     cell_mesh=df_vtk.dc[ds_id, 'cell_mesh', mamut_t].dc.path[0],\n",
    "#                     rgb_grid=df_vtk.dc[ds_id, 'rgb_grid', mamut_t].dc.path[0])\n",
    "\n",
    "# # db = OrganoidViewer()\n",
    "\n",
    "# panel = db.panel().servable()\n",
    "# panel\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# mamut_t = 421\n",
    "# ds_id = 2\n",
    "\n",
    "# db.load(features=dfn.set_index('mamut_t').loc[mamut_t:mamut_t],\n",
    "#         nuclei_mesh=df_vtk.dc[ds_id, 'nuclei_mesh', mamut_t].dc.path[0],\n",
    "#         cell_mesh=df_vtk.dc[ds_id, 'cell_mesh', mamut_t].dc.path[0],\n",
    "#         rgb_grid=df_vtk.dc[ds_id, 'rgb_grid', mamut_t].dc.path[0],\n",
    "#         features_bounds=features_bounds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dashboard\n",
    "Putting it all together. At the end it will start the interactive panel on the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from inter_view.io import CollectionHandler\n",
    "\n",
    "class OrganoidDashboard(param.Parameterized):\n",
    "    \n",
    "    df_vtk = param.DataFrame(pd.DataFrame())\n",
    "    df_h5 = param.DataFrame(pd.DataFrame())\n",
    "    \n",
    "    dataset = param.ObjectSelector(default=0, objects=[0])\n",
    "    timepoint = param.Integer(100, bounds=(1, 100))\n",
    "    \n",
    "    tree_viewer = param.Parameter(TreeViewer(node_color='nuclei_volume'), instantiate=False)\n",
    "    organoid_viewer = param.Parameter(OrganoidViewer(nuclei_color='nuclei_volume'))\n",
    "    \n",
    "    def __init__(self, *args, **kargs):\n",
    "        super().__init__(*args, **kargs)\n",
    "        \n",
    "        self.param.dataset.objects = self.df_h5.index.tolist()\n",
    "        if self.dataset not in self.param.dataset.objects:\n",
    "            self.dataset = self.param.dataset.objects[0]\n",
    "            \n",
    "        self.export_button = pn.widgets.FileDownload(label='export', callback=self._export_callback, filename='tmp.svg')\n",
    "            \n",
    "        self._update_tree_viewer()\n",
    "        self._update_sample()\n",
    "            \n",
    "    @param.depends('timepoint', watch=True)\n",
    "    def _update_sample(self):\n",
    "        dfn = self.tree_viewer.dfn\n",
    "        self.organoid_viewer.load(features=dfn.set_index('mamut_t').loc[self.timepoint:self.timepoint],\n",
    "                                  nuclei_mesh=self.df_vtk.dc[self.dataset, 'nuclei_mesh', self.timepoint].dc.path[0],\n",
    "                                  cell_mesh=self.df_vtk.dc[self.dataset, 'cell_mesh', self.timepoint].dc.path[0],\n",
    "                                  rgb_grid=self.df_vtk.dc[self.dataset, 'rgb_grid', self.timepoint].dc.path[0])\n",
    "        \n",
    "    @param.depends('dataset', watch=True)\n",
    "    def _update_tree_viewer(self):\n",
    "        self.tree_viewer.h5_path = self.df_h5.dc[self.dataset].dc.path[0]\n",
    "        \n",
    "        if self.tree_viewer.dfn is not None:\n",
    "            self._compute_features_bounds()\n",
    "            \n",
    "            t_bounds = tuple(self.tree_viewer.dfn.mamut_t.agg(['min', 'max']))\n",
    "            self.param.timepoint.bounds = t_bounds\n",
    "            self.timepoint = (t_bounds[0] + t_bounds[1])//2\n",
    "            self._update_sample()\n",
    "        else:\n",
    "            self.organoid_viewer.clear_renderer()\n",
    "        \n",
    "    def _compute_features_bounds(self):\n",
    "        self.organoid_viewer.features_bounds = self.tree_viewer.dfn.agg(['min', 'max']).to_dict('list')\n",
    "        self.organoid_viewer.features_bounds['selection'] = [0,1]\n",
    "        self.organoid_viewer.features_bounds['neighbor_trace'] = [0,2]\n",
    "    \n",
    "    @param.depends('tree_viewer.selection_update_counter', watch=True)\n",
    "    def _update_feature_cmap(self):\n",
    "        dfn = self.tree_viewer.dfn\n",
    "        features = dfn.set_index('mamut_t').loc[self.timepoint:self.timepoint]\n",
    "        self.organoid_viewer.update_features(features)\n",
    "    \n",
    "    @param.depends('tree_viewer.last_clicked_mamut_t', watch=True)\n",
    "    def _select_timepoint(self):\n",
    "        self.timepoint = self.tree_viewer.last_clicked_mamut_t\n",
    "    \n",
    "    def export_mesh(self):\n",
    "        code = '''\n",
    "                renderWindow.captureImages()[0].then(\n",
    "                  (image) => {{\n",
    "                    var download = document.createElement('a');\n",
    "                    download.href = image;\n",
    "                    download.download = 'DS{}_T{}_N-{}_C-{}_mesh.png';\n",
    "                    download.click();\n",
    "                  }}\n",
    "                );\n",
    "            '''\n",
    "        \n",
    "        code = code.format(self.dataset, self.timepoint, self.organoid_viewer.nuclei_color, self.organoid_viewer.cell_color)\n",
    "        \n",
    "        self.organoid_viewer.jscode.run(code)\n",
    "    \n",
    "    def export_plot(self):\n",
    "        try:\n",
    "            hv.extension('matplotlib', logo=False)\n",
    "\n",
    "            if self.tree_viewer.node_size == 'fixed':\n",
    "                node_size = 5\n",
    "            else:\n",
    "                node_size = self.tree_viewer.node_size\n",
    "\n",
    "            if self.tree_viewer.edge_size == 'fixed':\n",
    "                edge_size = 2\n",
    "            else:\n",
    "                edge_size = self.tree_viewer.edge_size\n",
    "\n",
    "            hv_tree = plot_tree(self.tree_viewer.dfn,\n",
    "                                 self.tree_viewer.dfe,\n",
    "                                 tooltips=[],\n",
    "                                 node_size=node_size * 10,\n",
    "                                 node_color=self.tree_viewer.node_color,\n",
    "                                 edge_width=edge_size,\n",
    "                                 edge_color=self.tree_viewer.edge_color,\n",
    "                                 min_node_size=2,\n",
    "                                 max_node_size=30,\n",
    "                                 min_edge_width=1,\n",
    "                                 max_edge_width=9,\n",
    "                                 backend='matplotlib')\n",
    "\n",
    "            hv_tree.opts(opts.Overlay('tree', show_frame=False, yaxis=None, show_title=False, aspect=0.7, framewise=True),\n",
    "                         opts.Points('tree', cmap=clipped_plasma_r, framewise=True),\n",
    "                         opts.Segments('tree', color='black', cmap=clipped_plasma_r, framewise=True))\n",
    "\n",
    "            layout = [hv_tree.opts(sublabel_size=0)]\n",
    "\n",
    "            cbar_img = np.linspace(0, 1, 64)[:,None]\n",
    "            if self.tree_viewer.node_color is not None:\n",
    "                cbar = plot_colorbar('clipped_plasma_r', self.organoid_viewer.features_bounds[self.tree_viewer.node_color])\n",
    "                layout.append(cbar)\n",
    "\n",
    "            if self.organoid_viewer.nuclei_color is not None:\n",
    "                cbar = plot_colorbar('clipped_plasma_r', self.organoid_viewer.features_bounds[self.organoid_viewer.nuclei_color])\n",
    "                layout.append(cbar)\n",
    "\n",
    "            if self.organoid_viewer.cell_color is not None:\n",
    "                cbar = plot_colorbar('clipped_plasma_r', self.organoid_viewer.features_bounds[self.organoid_viewer.cell_color])\n",
    "                layout.append(cbar)\n",
    "\n",
    "            buff = BytesIO()\n",
    "\n",
    "            hv.save(hv.Layout(layout).opts(framewise=True, fig_size=800, tight=True, aspect_weight=True, axiswise=True), buff, fmt='pdf')\n",
    "            buff.seek(0)\n",
    "            hv.extension('bokeh', logo=False)\n",
    "            \n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            hv.extension('bokeh', logo=False)\n",
    "        \n",
    "        return buff\n",
    "    \n",
    "    def _export_callback(self):\n",
    "        self.export_mesh()\n",
    "        buff = self.export_plot()\n",
    "        self.export_button.filename = 'DS{}_color-{}_tree.pdf'.format(self.dataset,\n",
    "                                                                      self.tree_viewer.node_color)\n",
    "        return buff\n",
    "        \n",
    "    def panel(self):\n",
    "        timepoint_vline = hv.VLine(x=self.tree_viewer.time_interval * self.timepoint)\n",
    "        \n",
    "        plots = pn.Row(timepoint_vline * self.tree_viewer.plot(),\n",
    "                       self.organoid_viewer.vtkpan,\n",
    "                       sizing_mode='stretch_both')\n",
    "        \n",
    "        widgets = pn.Row(pn.Column(pn.WidgetBox(self.param.dataset,\n",
    "                                                self.param.timepoint),\n",
    "                                   self.export_button,\n",
    "                                   # TODO remove dummy button when panel bug fixed https://github.com/holoviz/panel/issues/1920\n",
    "                                   pn.widgets.Button(width=0, margin=0)),\n",
    "                         self.tree_viewer.widgets(),\n",
    "                         self.organoid_viewer.widgets())\n",
    "        \n",
    "        # js link timepoint slider to vertical line\n",
    "        code = '''glyph.location = source.value * {}'''.format(self.tree_viewer.time_interval)\n",
    "        widgets[0][0][1].jslink(timepoint_vline, code={'value': code})\n",
    "\n",
    "        return pn.Column(plots, widgets, sizing_mode='stretch_both').servable()\n",
    "\n",
    "db = OrganoidDashboard(df_h5=df_h5,\n",
    "                       df_vtk=df_vtk)\n",
    "\n",
    "panel = db.panel()\n",
    "panel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# First neighbors\n",
    "Initial idea of how to calculate first neighbors. Keeping in mind that currently the dashboard does not allow the back-and-forth transfer of information in an easy way, such that selecting a node for further calculation of not yet existing feature is currently not possible in an interactive way. For the time being the way around this is to run the cells below first and then rerun the dashboard."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trace_back_neighbors(tree, node_id):\n",
    "    # TODO check that delaunay_neighbors exist\n",
    "    \n",
    "    tree.set_all_nodes_attribute('neighbor_trace', 0)\n",
    "    \n",
    "    def label_ancestors(node_id):\n",
    "        current_node = tree.nodes[node_id]\n",
    "        current_node['neighbor_trace'] = 2\n",
    "        \n",
    "        neighbors_ids = current_node['cell_neighbors']\n",
    "        if isinstance(neighbors_ids, list):\n",
    "            for nb_id in neighbors_ids:\n",
    "                tree.nodes[nb_id]['neighbor_trace'] = 1\n",
    "\n",
    "        for p_id in tree.predecessors(node_id):\n",
    "            label_ancestors(p_id)\n",
    "\n",
    "    label_ancestors(node_id)\n",
    "    \n",
    "    \n",
    "from lstree.lineage.plot import tree_to_dataframe, plot_tree\n",
    "from lstree.lineage.ltree import LineageTree\n",
    "\n",
    "# needed once per dataset\n",
    "dfn = db.tree_viewer.dfn\n",
    "dfe = db.tree_viewer.dfe\n",
    "tree = LineageTree.from_dataframe(dfn, dfe, 'mamut_t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace_back_neighbors(tree, 1774)\n",
    "\n",
    "dfnb, dfeb = tree_to_dataframe(tree)\n",
    "dfn['neighbor_trace'] = dfnb['neighbor_trace']\n",
    "\n",
    "db.tree_viewer.dfn = dfn\n",
    "db.tree_viewer.tree_update_counter += 1\n",
    "db._update_feature_cmap()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}