{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import numpy as np \n",
    "from brainlit.preprocessing import image_process\n",
    "from brainlit.algorithms.connect_fragments import dynamic_programming_viterbi2\n",
    "import napari\n",
    "import networkx as nx\n",
    "from napari.utils import nbscreenshot\n",
    "\n",
    "%gui qt"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Viterbi2 Demo"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Generate basic image and labels"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "image = 0.5*np.ones((100,100,1))\n",
    "image[50:55,0:25,0] = 0.91\n",
    "image[50:55,30:50,0] = 0.92\n",
    "image[45:50,55:75,0] = 0.93\n",
    "image[60:65,55:75,0] = 0.94\n",
    "image[45:60,85:,0] = 0.95\n",
    "\n",
    "labels = np.zeros((100,100,1), dtype=int)\n",
    "labels[50:55,0:25,0] = 1\n",
    "labels[50:55,30:50,0] = 2\n",
    "labels[45:50,55:75,0] = 3\n",
    "labels[60:65,55:75,0] = 4\n",
    "labels[45:60,85:,0] = 5\n",
    "\n",
    "axon_coords = [[52, 2, 0]]\n",
    "soma_coords = [[50, 90, 0]]\n",
    "\n",
    "res = [0.1,0.1,0.1]\n",
    "\n",
    "# viewer = napari.Viewer(ndisplay=3)\n",
    "# viewer.add_image(image)\n",
    "# viewer.add_labels(labels)\n",
    "# viewer.camera.angles = [0, -90, 180]\n",
    "# nbscreenshot(viewer)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Find labels of coordinates"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "_, axon_lbls = image_process.label_points(labels, axon_coords, res)\n",
    "axon_lbl= axon_lbls[0]\n",
    "_, soma_lbls = image_process.label_points(labels, soma_coords, res)\n",
    "soma_lbl = soma_lbls[0]"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Create Reconstruction Object"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "mpnp = dynamic_programming_viterbi2.most_probable_neuron_path(image, labels, [soma_lbl], res)\n",
    "mpnp.frags_to_lines()\n",
    "mpnp.compute_all_costs_dist(point_point_func = mpnp.point_point_dist, point_blob_func = mpnp.point_blob_dist)\n",
    "mpnp.compute_all_costs_int()\n",
    "mpnp.create_nx_graph()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Run Reconstruction"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "path_states = nx.shortest_path(mpnp.nxGraph, 0, 8, weight='weight')\n",
    "path_comps = [mpnp.state_to_comp[state][1] for state in path_states]\n",
    "print(path_comps)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Plot Reconstruction"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# path_comps = []\n",
    "# for state in path_states:\n",
    "#     path_comps.append(mpnp.state_to_comp[state][1])\n",
    "# print(f'path sequence: {path_states}')\n",
    "# print(f'component sequence: {path_comps}')\n",
    "\n",
    "# path_mask = 0*labels\n",
    "# for i, label in enumerate(path_comps):\n",
    "#     path_mask[labels == label] = i+1\n",
    "\n",
    "# soma_mask = 0*labels\n",
    "# for soma_lbl in mpnp.soma_lbls:\n",
    "#     soma_mask[labels == soma_lbl] = soma_lbl\n",
    "\n",
    "# viewer = napari.Viewer(ndisplay=3)\n",
    "# viewer.add_image(mpnp.image)\n",
    "# viewer.add_labels(labels)\n",
    "# viewer.add_labels(path_mask)\n",
    "# viewer.add_labels(soma_mask)\n",
    "# viewer.add_labels(labels == axon_lbl)\n",
    "\n",
    "# viewer.add_points([axon_coords[0]], face_color=\"red\", size=10)\n",
    "\n",
    "# lines = []\n",
    "# cumul_cost = 0\n",
    "# for s, state in enumerate(path_states):\n",
    "#     if s>0:\n",
    "#         dist_cost = mpnp.cost_mat_dist[path_states[s-1], state]\n",
    "#         int_cost = mpnp.cost_mat_int[path_states[s-1], state]\n",
    "#         cumul_cost += dist_cost + int_cost\n",
    "#         print(f\"Trans. #{s}: dist cost state {path_states[s-1]}->state {state}, comp {mpnp.state_to_comp[path_states[s-1]][1]}->comp {mpnp.state_to_comp[state][1]}: {dist_cost:.2f}, int cost: {int_cost:.2f}, cum. cost: {cumul_cost:.2f}\")\n",
    "#     if mpnp.state_to_comp[state][0] == \"fragment\":\n",
    "#         lines.append(list(mpnp.state_to_comp[state][2][\"coord1\"]))\n",
    "#         lines.append(list(mpnp.state_to_comp[state][2][\"coord2\"]))\n",
    "#     elif mpnp.state_to_comp[path_states[s-1]][0] == \"fragment\":\n",
    "#         lines.append(list(mpnp.state_to_comp[path_states[s-1]][2][\"soma connection point\"]))     \n",
    "# viewer.add_shapes(lines, shape_type=\"path\", edge_color=\"blue\", edge_width=2)\n",
    "\n",
    "# viewer.camera.angles = [0, -90, 180]\n",
    "# nbscreenshot(viewer)\n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "orig_nbformat": 4,
  "language_info": {
   "name": "python",
   "version": "3.8.10",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.10 64-bit ('docs_env': venv)"
  },
  "interpreter": {
   "hash": "c6a82fd7624a30cb39f184f8a867df460926136b3ed0e9f03cd044bdf3194e37"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}