{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exploring `doit` graphs with `ipywidgets`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as W, traitlets as T, doit, graphviz, jinja2, textwrap, re; ui = W.VBox([W.HTML('restart and run all to see something here')]); ui"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Doit(W.Widget):\n",
    "    doit_ = T.Instance(doit.doit_cmd.DoitMain)\n",
    "    tasks = W.trait_types.TypedTuple(trait=T.Instance(doit.task.Task))\n",
    "    \n",
    "    @T.observe(\"doit_\")\n",
    "    def on_doit_change(self, change):\n",
    "        cmds = self.doit_.get_cmds()\n",
    "        tasks, _ = self.doit_.task_loader.load_tasks(cmds['list'], None, None)\n",
    "        tc = doit.control.TaskControl(tasks)\n",
    "        self.tasks = tuple(tc.tasks.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FileDoit(Doit):\n",
    "    path = T.Unicode(\"dodo.py\").tag(sync=True)\n",
    "    \n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.on_path_change(T.Bunch(new=self.path))\n",
    "    \n",
    "    @T.observe(\"path\")\n",
    "    def _on_path_change(self, change):\n",
    "        self.on_path_change(change)    \n",
    "    \n",
    "    def on_path_change(self, change):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PyFileDoit(FileDoit):    \n",
    "    def on_path_change(self, change=None):\n",
    "        self.doit_ = doit.doit_cmd.DoitMain(doit.cmd_base.ModuleTaskLoader(doit.loader.get_module(self.path)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "it = PyFileDoit(path = \"dodo.py\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TaskGraph(W.HTML):\n",
    "    tasks = W.trait_types.TypedTuple(trait=T.Instance(doit.task.Task))\n",
    "    graph = T.Instance(graphviz.Digraph)\n",
    "    node_template = T.Unicode(\"\"\"<\n",
    "    <TABLE BORDER=\"0\" CELLBORDER=\"0\" CELLSPACING=\"0\"><TR><TD><b>{{t.name}}</b></TD></TR><TR><TD>{{wrap(t.doc)}}</TD></TR></TABLE>\n",
    "    >\"\"\")\n",
    "    filter_text = T.Unicode()\n",
    "    filter_kind = T.Unicode()\n",
    "    rankdir = T.Unicode(\"RL\")\n",
    "    wrap = T.Int(24)\n",
    "    \n",
    "    @T.observe(\"tasks\", \"node_template\", \"filter_text\", \"filter_kind\", \"wrap\", \"rankdir\")\n",
    "    def on_tasks(self, change):\n",
    "        tmpl = jinja2.Template(self.node_template)\n",
    "        wrap = lambda x: \"<br/>\".join(textwrap.wrap(x, self.wrap))\n",
    "        self.graph = G = graphviz.Digraph(format=\"svg\", node_attr=dict(fontname=\"sans-serif\"))\n",
    "        G.attr(rankdir=self.rankdir)\n",
    "        tasks = self.filtered()\n",
    "        task_names = [t.name for t in tasks]\n",
    "        for t in tasks:\n",
    "            G.node(t.name, tmpl.render(t=t, wrap=wrap), shape=\"none\")\n",
    "            for tdep in t.task_dep:\n",
    "                if tdep in task_names:\n",
    "                    G.edge(tdep, t.name)\n",
    "        \n",
    "        self.value = re.sub(r'<svg (.*)viewBox', '<svg viewBox', G._repr_svg_(), flags=re.M | re.DOTALL)\n",
    "    \n",
    "    def filtered(self):\n",
    "        tasks = self.tasks or []\n",
    "        kind = self.filter_kind\n",
    "        if not kind:\n",
    "            return tasks\n",
    "        \n",
    "        if kind in [\"upstream\", \"downstream\"]:\n",
    "            filtered = {t.name: t for t in tasks if re.findall(self.filter_text, t.name)}\n",
    "        elif kind == \"file\":\n",
    "            filtered = {t.name: t for t in tasks if [fd for fd in t.file_dep if re.findall(self.filter_text, fd)]}\n",
    "        \n",
    "        candidates = {t.name: t for t in tasks if t.name not in filtered}\n",
    "        \n",
    "        while True:\n",
    "            found = False\n",
    "            f_round = list(filtered.values())\n",
    "            for t in f_round:\n",
    "                if kind == \"upstream\":\n",
    "                    for d in t.task_dep:\n",
    "                        if d in candidates:\n",
    "                            filtered[d] = candidates.pop(d)\n",
    "                            found = True\n",
    "                elif kind in [\"downstream\", \"file\"]:\n",
    "                    d_round = list(candidates.values())\n",
    "                    for d in d_round:\n",
    "                        if t.name in d.task_dep:\n",
    "                            filtered[d.name] = candidates.pop(d.name)\n",
    "                            found = True\n",
    "            if not found:\n",
    "                break\n",
    "\n",
    "        return filtered.values()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tg = TaskGraph()\n",
    "wrap = W.IntSlider(24, min=5, max=45, description=\"wrap\")\n",
    "filter_text = W.Text(\".*\", description=\"goal\")\n",
    "filter_kind = W.SelectionSlider(options=[\"upstream\", \"downstream\", \"file\", \"\"], description=\"filter by\")\n",
    "rankdir = W.SelectionSlider(options=[\"LR\", \"TB\"], description=\"direction\")\n",
    "\n",
    "T.link((it, \"tasks\"), (tg, \"tasks\"))\n",
    "T.link((filter_text, \"value\"), (tg, \"filter_text\"))\n",
    "T.link((wrap, \"value\"), (tg, \"wrap\"))\n",
    "T.link((filter_kind, \"value\"), (tg, \"filter_kind\"))\n",
    "T.link((rankdir, \"value\"), (tg, \"rankdir\"))\n",
    "\n",
    "ui.children = [W.HBox([filter_text, filter_kind, rankdir, wrap]), tg]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}