{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C_h806vN-uqg"
   },
   "source": [
    "# Programmatically Interacting with the Fish1 Dataset via CAVE\n",
    "\n",
    "This notebook is a getting-started guide for accessing the **Fish1 whole-brain connectomics dataset** through the [CAVE (Connectome Annotation and Versioning Engine)](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC10402030/) Python client.\n",
    "\n",
    "By the end you will be able to:\n",
    "1. Authenticate and connect to the Fish1 CAVE server\n",
    "2. Explore the annotation tables (somas, synapses, cell types)\n",
    "3. Look up a neuron by coordinate or by its stable lore ID\n",
    "4. Query incoming / outgoing synapses and filter them by type or size\n",
    "5. Export a skeleton as an SWC file\n",
    "7. List all proofread (edited) segments\n",
    "\n",
    "**ID system quick reference**\n",
    "\n",
    "| Term | Description |\n",
    "|------|-------------|\n",
    "| **lore ID** | Small stable integer identifying a soma (e.g. `173502`). Does **not** change after proofreading. |\n",
    "| **root ID** | Large 64-bit integer for the current segmentation of a neuron (e.g. `864691128630286274`). Changes after merge / split edits \u2014 always check `is_latest_roots()` before use. |\n",
    "| **supervoxel ID** | Leaf-level voxel chunk; the atomic unit of the segmentation graph. |\n"
   ],
   "id": "C_h806vN-uqg"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bMyRyx8T-uqj"
   },
   "source": [
    "---\n",
    "## 1 \u00b7 Installation & Imports\n"
   ],
   "id": "bMyRyx8T-uqj"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KVAdbuqz-uqj",
    "outputId": "c686d569-5e7f-4e98-d9b9-8de20bc82771"
   },
   "outputs": [],
   "source": [
    "# Run once (comment out after first execution)\n",
    "!pip install caveclient cloud-volume"
   ],
   "id": "KVAdbuqz-uqj"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "T1_rucSk-uqk"
   },
   "outputs": [],
   "source": [
    "import caveclient\n",
    "from caveclient import CAVEclient, chunkedgraph as cg\n",
    "from cloudvolume import CloudVolume\n",
    "import numpy as np\n",
    "import pandas as pd"
   ],
   "id": "T1_rucSk-uqk"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cYpRJ5DG-uql"
   },
   "source": "# Paste your token below and run this cell\nauth.save_token(token=\"YOUR_TOKEN_HERE\", overwrite=True)",
   "id": "cYpRJ5DG-uql"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4DL9TDGn-uqm",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "5cf78379-3057-4f97-b5ba-2aa87003051c"
   },
   "outputs": [],
   "source": [
    "global_url = \"https://global.brain-wire-test.org/\"\n",
    "auth = caveclient.auth.AuthClient(server_address=global_url)\n",
    "\n",
    "# Opens a browser page \u2014 copy the token string shown there\n",
    "auth.setup_token(make_new=False)"
   ],
   "id": "4DL9TDGn-uqm"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pu81jS15-uqn"
   },
   "outputs": [],
   "source": "# Paste your token below and run this cell\nauth.save_token(token=\"YOUR_TOKEN_HERE\", overwrite=True)",
   "id": "pu81jS15-uqn"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pKxZhnv4-uqo"
   },
   "source": [
    "---\n",
    "## 3 \u00b7 Connect to the Fish1 Server\n",
    "\n",
    "Three clients are used throughout this notebook:\n",
    "\n",
    "| Client | Purpose |\n",
    "|--------|---------|\n",
    "| `client` | Main CAVE client \u2014 annotation tables, materialization, auth |\n",
    "| `cggraph` | ChunkedGraph client \u2014 root-ID lookups, edit history |\n",
    "| `cv` | CloudVolume \u2014 segmentation voxel data (coordinate \u2192 supervoxel) |\n"
   ],
   "id": "pKxZhnv4-uqo"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "g95Zzss8-uqo",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "fb834aea-fe9e-445a-df2a-8e24e258c581"
   },
   "outputs": [],
   "source": [
    "global_url = \"https://global.brain-wire-test.org/\"\n",
    "local_url  = \"https://pcgv3local.brain-wire-test.org\"\n",
    "datastack  = \"fish1_full\"\n",
    "dataset    = \"fish1_v250915\"\n",
    "\n",
    "# Main annotation / materialization client\n",
    "client = CAVEclient(datastack_name=datastack, server_address=global_url)\n",
    "client.materialize.version = client.materialize.most_recent_version()\n",
    "print(f\"Using materialization version: {client.materialize.version}\")\n",
    "\n",
    "# ChunkedGraph client (root-ID resolution, edit history)\n",
    "cggraph = cg.ChunkedGraphClient(\n",
    "    server_address=local_url,\n",
    "    table_name=dataset,\n",
    "    auth_client=caveclient.auth.AuthClient(token=client.auth.token),\n",
    ")\n",
    "\n",
    "# Segmentation volume (coordinate \u2192 supervoxel lookups)\n",
    "cv = client.info.segmentation_cloudvolume()\n",
    "\n",
    "print(\"Connected successfully.\")"
   ],
   "id": "g95Zzss8-uqo"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HgaXadKV-uqo"
   },
   "source": [
    "---\n",
    "## 4 \u00b7 Exploring Annotation Tables\n",
    "\n",
    "All annotations are stored in versioned, queryable tables accessed via `client.materialize`."
   ],
   "id": "HgaXadKV-uqo"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "g0IBHZYq-uqp",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "e6b12ab5-3cc8-4a89-9a84-d322e26f9956"
   },
   "outputs": [],
   "source": [
    "#@title List all available tables\n",
    "all_tables = client.annotation.get_tables()\n",
    "print(all_tables)"
   ],
   "id": "g0IBHZYq-uqp"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EMYe3V3r-uqp",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "1ba116f2-be71-4067-fd79-652dfaa9c945"
   },
   "outputs": [],
   "source": [
    "#@title Inspect metadata for a table (change the name to explore others)\n",
    "client.annotation.get_table_metadata(\"somas\")"
   ],
   "id": "EMYe3V3r-uqp"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7nmxPIl6-uqp",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 432
    },
    "outputId": "630d7a75-85b5-4802-b3ba-63eece4fdcad"
   },
   "outputs": [],
   "source": [
    "#@title Preview the somas table\n",
    "# Column reference:\n",
    "#   id               \u2014 stable lore ID (soma's fixed identifier)\n",
    "#   cell_type        \u2014 \"exc\" (excitatory), \"inh\" (inhibitory), or \"na\"\n",
    "#   pt_root_id       \u2014 current root ID for this segment\n",
    "#   pt_supervoxel_id \u2014 supervoxel at the soma location\n",
    "#   pt_position      \u2014 [x, y, z] voxel coords at 16 nm \u00d7 16 nm \u00d7 30 nm\n",
    "#   created          \u2014 timestamp when annotation was added\n",
    "client.materialize.tables.somas().query(limit=10)"
   ],
   "id": "7nmxPIl6-uqp"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fk112IPi-uqp"
   },
   "source": [
    "---\n",
    "## 5 \u00b7 Looking Up Neurons\n",
    "\n",
    "### 5a \u00b7 By 3-D coordinate\n",
    "\n",
    "Given a voxel coordinate (e.g. from Neuroglancer), retrieve the supervoxel and root ID, then confirm it is the most recent root."
   ],
   "id": "fk112IPi-uqp"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "W2Nt-KzH-uqp",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "cc0bfc83-ce38-4a45-de8a-1aee9c186227"
   },
   "outputs": [],
   "source": [
    "#@title Look up a neuron by voxel coordinate\n",
    "def get_latest_root_id(x, y, z, resolution=(8, 8, 30)):\n",
    "    \"\"\"Return (supervoxel_id, root_id) for a voxel coordinate.\n",
    "\n",
    "    Args:\n",
    "        x, y, z:    voxel coordinates\n",
    "        resolution: voxel size in nm (default 8, 8, 30)\n",
    "\n",
    "    Returns:\n",
    "        (supervoxel_id, latest_root_id)\n",
    "    \"\"\"\n",
    "    supervoxel_id = cv.download_point(\n",
    "        pt=(x, y, z), size=1, agglomerate=False, coord_resolution=resolution\n",
    "    )\n",
    "    supervoxel_id = np.int64(supervoxel_id[0, 0, 0, 0])\n",
    "    print(f\"supervoxel_id : {supervoxel_id}\")\n",
    "\n",
    "    root_id = cggraph.get_root_id(supervoxel_id)\n",
    "    print(f\"root_id       : {root_id}\")\n",
    "\n",
    "    if not cggraph.is_latest_roots([root_id]):\n",
    "        root_id = cggraph.get_latest_roots(root_id)\n",
    "        print(f\"updated root  : {root_id}  (segment was edited)\")\n",
    "    else:\n",
    "        print(\"root_id is current\")\n",
    "\n",
    "    return supervoxel_id, root_id\n",
    "\n",
    "\n",
    "# Example \u2014 a hindbrain integrator neuron\n",
    "supervoxel_id, root_id = get_latest_root_id(75765, 38366, 6682)"
   ],
   "id": "W2Nt-KzH-uqp"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DR7gO3ko-uqq"
   },
   "source": [
    "### 5b \u00b7 By lore ID\n",
    "\n",
    "The lore ID is the stable soma identifier. Use it to retrieve the current root ID, voxel position, and excitatory/inhibitory cell type."
   ],
   "id": "DR7gO3ko-uqq"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JRGGouhn-uqq",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "a98373eb-1af7-4593-f2dc-b7ccf2a2b530"
   },
   "outputs": [],
   "source": [
    "#@title Look up a neuron by lore ID\n",
    "lore_id = \"173502\"  # \u2190 change this to query a different neuron\n",
    "\n",
    "row = client.materialize.tables.somas(id=lore_id).query().iloc[0]\n",
    "\n",
    "root_id   = int(row[\"pt_root_id\"])\n",
    "sv_id     = int(row[\"pt_supervoxel_id\"])\n",
    "pt_vox    = row[\"pt_position\"]\n",
    "cell_type = row[\"cell_type\"]\n",
    "\n",
    "type_label = {\"exc\": \"excitatory\", \"inh\": \"inhibitory\"}.get(cell_type, \"unknown\")\n",
    "\n",
    "print(f\"Lore ID   : {lore_id}\")\n",
    "print(f\"Root ID   : {root_id}\")\n",
    "print(f\"Position  : {pt_vox}  (voxels @ 16\u00d716\u00d730 nm)\")\n",
    "print(f\"Cell type : {type_label}\")"
   ],
   "id": "JRGGouhn-uqq"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eWpvG5O0-uqq"
   },
   "source": [
    "### 5c \u00b7 Root ID \u2192 lore ID (reverse lookup)"
   ],
   "id": "eWpvG5O0-uqq"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mDCMyd_m-uqq",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "3aeb0e7b-6644-48d8-e413-2d16fdf9a3a4"
   },
   "outputs": [],
   "source": [
    "#@title Reverse lookup: root ID \u2192 lore ID\n",
    "# Uses root_id defined in 5a or 5b above\n",
    "row = client.materialize.tables.somas(pt_root_id=root_id).query()\n",
    "\n",
    "if len(row):\n",
    "    print(f\"Root {root_id} \u2192 lore ID {row.iloc[0]['id']}\")\n",
    "else:\n",
    "    print(f\"No soma annotation found for root {root_id}\")"
   ],
   "id": "mDCMyd_m-uqq"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UdhGw8U1-uqq"
   },
   "source": [
    "### 5d \u00b7 Get all supervoxels for a root ID\n",
    "\n",
    "Useful when you need every leaf-level segment belonging to a neuron, e.g. for volume calculations."
   ],
   "id": "UdhGw8U1-uqq"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dJo2UjcJ-uqq",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "57249bcd-2188-4742-df95-1a64eed7f6b9"
   },
   "outputs": [],
   "source": [
    "#@title List all supervoxels for a root ID\n",
    "def get_supervoxels(root_id, client):\n",
    "    \"\"\"Return all supervoxel IDs (layer-1 leaves) for a root ID.\"\"\"\n",
    "    return client.chunkedgraph.get_leaves(root_id, stop_layer=1)\n",
    "\n",
    "svids = get_supervoxels(root_id, client)\n",
    "print(f\"Root {root_id} \u2192 {len(svids)} supervoxels\")"
   ],
   "id": "dJo2UjcJ-uqq"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Wk6NDbmH-uqq"
   },
   "source": [
    "---\n",
    "## 6 \u00b7 Querying Synapses\n",
    "\n",
    "Synapse tables encode axon-to-dendrite connections with:\n",
    "- **pre_pt_root_id** / **post_pt_root_id** \u2014 partner root IDs\n",
    "- **pre_pt_position** / **post_pt_position** \u2014 spatial locations\n",
    "- **tag** \u2014 `'1'` = inhibitory, `'2'` = excitatory (from the `synapses_axde_label` table)\n",
    "\n",
    "The example below uses `lore_id = \"173502\"`. Change it to analyse any neuron."
   ],
   "id": "Wk6NDbmH-uqq"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8F_y0TF8-uqq",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "46fe3556-441e-4bd6-ac2d-d23982dc3e03"
   },
   "outputs": [],
   "source": [
    "#@title Fetch incoming and outgoing synapses for a neuron\n",
    "lore_id = \"173502\"  # \u2190 change this\n",
    "\n",
    "row     = client.materialize.tables.somas(id=lore_id).query().iloc[0]\n",
    "root_id = int(row[\"pt_root_id\"])\n",
    "print(f\"Lore ID {lore_id} \u2192 root ID {root_id}\\n\")\n",
    "\n",
    "# Incoming (post-synaptic)\n",
    "incoming = client.materialize.tables.synapses_axde_label(post_pt_root_id=root_id).query()\n",
    "inh_in   = incoming[incoming[\"tag\"] == \"1\"]\n",
    "exc_in   = incoming[incoming[\"tag\"] == \"2\"]\n",
    "\n",
    "# Outgoing (pre-synaptic)\n",
    "outgoing = client.materialize.tables.synapses_axde_label(pre_pt_root_id=root_id).query()\n",
    "\n",
    "print(f\"Incoming synapses : {len(incoming):>5}  \"\n",
    "      f\"({len(inh_in)} inhibitory, {len(exc_in)} excitatory)\")\n",
    "print(f\"  unique pre-synaptic partners : {incoming['pre_pt_root_id'].nunique()}\")\n",
    "print(f\"Outgoing synapses : {len(outgoing):>5}\")"
   ],
   "id": "8F_y0TF8-uqq"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F2nDUHc--uqr"
   },
   "source": [
    "### 6b \u00b7 Assemble a full synapse table with size IDs and e/i labels\n",
    "\n",
    "The full synapse record (including pre- and post-synaptic site IDs) is spread across three tables. The cell below joins them."
   ],
   "id": "F2nDUHc--uqr"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CDIaBsoQ-uqr",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 223
    },
    "outputId": "40860cfe-7918-4254-fc91-0ee05b18f66a"
   },
   "outputs": [],
   "source": [
    "#@title Build a merged synapse table\n",
    "# Fetch size/ID annotations for each synapse\n",
    "pre_ids  = client.materialize.tables.synapses_axde_pre_synapse_id(\n",
    "               id=incoming.id.values.tolist()).query()[[\"id\", \"tag\"]]\n",
    "post_ids = client.materialize.tables.synapses_axde_post_synapse_id(\n",
    "               id=incoming.id.values.tolist()).query()[[\"id\", \"tag\"]]\n",
    "\n",
    "merged = (\n",
    "    incoming\n",
    "    .merge(pre_ids,  on=\"id\", suffixes=(\"\", \"_pre\"))\n",
    "    .merge(post_ids, on=\"id\", suffixes=(\"\", \"_post\"))\n",
    ")\n",
    "\n",
    "synapse_table = merged[\n",
    "    [\"pre_pt_position\", \"tag_pre\", \"post_pt_position\", \"tag_post\", \"tag\"]\n",
    "].rename(columns={\n",
    "    \"tag_pre\":  \"pre_site_id\",\n",
    "    \"tag_post\": \"post_site_id\",\n",
    "    \"tag\":      \"ei_label\",      # '1'=inh, '2'=exc\n",
    "})\n",
    "\n",
    "print(f\"Merged table: {len(synapse_table)} rows\")\n",
    "synapse_table.head()"
   ],
   "id": "CDIaBsoQ-uqr"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qpWGD3f7-uqr",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "f9d7d796-7bac-4262-ca84-8ba223ee7f07"
   },
   "outputs": [],
   "source": [
    "#@title Export merged synapse table to CSV\n",
    "synapse_table.to_csv(f\"synapses_{lore_id}_{root_id}.csv\", index=False, header=True)\n",
    "print(f\"Saved synapses_{lore_id}_{root_id}.csv\")"
   ],
   "id": "qpWGD3f7-uqr"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GzHjrN1E-uqr"
   },
   "source": [
    "### 6c \u00b7 Filter synapses by physical size\n",
    "\n",
    "Each synapse has a bounding box stored in `synapses_axon_to_dendrite_size`. Use this to remove small / spurious contacts."
   ],
   "id": "GzHjrN1E-uqr"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QNceCs31-uqr",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "c0f7fb1e-9a51-411b-fa88-7bcbc0557276"
   },
   "outputs": [],
   "source": [
    "#@title Filter synapses by bounding-box size\n",
    "syn_sizes = client.materialize.tables.synapses_axon_to_dendrite_size(\n",
    "    tag=pre_ids.tag.values.tolist(),\n",
    "    tag2=post_ids.tag.values.tolist(),\n",
    ").query()\n",
    "\n",
    "# Compute bounding-box dimensions (voxels)\n",
    "bbox_start = np.vstack(syn_sizes[\"pt_position\"].values)\n",
    "bbox_end   = np.vstack(syn_sizes[\"pt2_position\"].values)\n",
    "dims       = np.abs(bbox_start - bbox_end)\n",
    "syn_sizes[[\"size_x\", \"size_y\", \"size_z\"]] = dims\n",
    "syn_sizes[\"volume\"] = dims[:, 0] * dims[:, 1] * dims[:, 2]\n",
    "\n",
    "# --- Filter options ---\n",
    "volume_threshold    = 1000          # voxels\u00b3\n",
    "dimension_threshold = (10, 10, 10)  # voxels per axis\n",
    "\n",
    "by_volume = syn_sizes[syn_sizes[\"volume\"] > volume_threshold]\n",
    "by_dims   = syn_sizes[\n",
    "    (syn_sizes[\"size_x\"] > dimension_threshold[0]) &\n",
    "    (syn_sizes[\"size_y\"] > dimension_threshold[1]) &\n",
    "    (syn_sizes[\"size_z\"] > dimension_threshold[2])\n",
    "]\n",
    "\n",
    "print(f\"All synapses         : {len(syn_sizes)}\")\n",
    "print(f\"After volume filter  : {len(by_volume)}\")\n",
    "print(f\"After per-axis filter: {len(by_dims)}\")"
   ],
   "id": "QNceCs31-uqr"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ufduK-_E-uqr"
   },
   "source": [
    "---\n",
    "## 7 \u00b7 Exporting Skeletons\n",
    "\n",
    "Skeletons are available in [SWC format](https://www.neuron.yale.edu/neuron/static/docs/neuronpython/swcformat.html) and can be loaded into tools such as [navis](https://navis.readthedocs.io), [NEURON](https://neuron.yale.edu), or [neuromorpho.org](https://neuromorpho.org)."
   ],
   "id": "ufduK-_E-uqr"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VvMsvOr8-uqr",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "d9d12adf-1e66-43c6-befe-dabe23841881"
   },
   "outputs": [],
   "source": [
    "#@title Download skeleton as SWC file\n",
    "lore_id = \"173502\"  # \u2190 change this\n",
    "\n",
    "row     = client.materialize.tables.somas(id=lore_id).query().iloc[0]\n",
    "root_id = int(row[\"pt_root_id\"])\n",
    "\n",
    "skel_vol = CloudVolume(\n",
    "    \"precomputed://middleauth+\"\n",
    "    \"https://pcgv3local.brain-wire-test.org\"\n",
    "    \"/skeletoncache/api/v1/fish1_full/precomputed/\"\n",
    ")\n",
    "skeleton = skel_vol.skeleton.get(root_id)\n",
    "\n",
    "out_file = f\"skeleton_{lore_id}_{root_id}.swc\"\n",
    "with open(out_file, \"w\") as f:\n",
    "    f.write(skeleton.to_swc())\n",
    "print(f\"Saved: {out_file}\")\n",
    "print(skeleton)  # summary of vertices / edges"
   ],
   "id": "VvMsvOr8-uqr"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4q-sNfDf-uqr",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "b07c1986-305e-40be-802b-e08756243731"
   },
   "outputs": [],
   "source": [
    "#@title Export skeleton + labelled synapses together\n",
    "synapses = client.materialize.tables.synapses_axde_label(post_pt_root_id=root_id).query()\n",
    "inh_syn  = synapses[synapses[\"tag\"] == \"1\"]\n",
    "exc_syn  = synapses[synapses[\"tag\"] == \"2\"]\n",
    "\n",
    "inh_syn.to_csv(f\"inh_synapses_{lore_id}_{root_id}.csv\", index=False)\n",
    "exc_syn.to_csv(f\"exc_synapses_{lore_id}_{root_id}.csv\", index=False)\n",
    "\n",
    "print(f\"Lore {lore_id}: {len(synapses)} total synapses  \"\n",
    "      f\"({len(inh_syn)} inhibitory, {len(exc_syn)} excitatory)\")\n",
    "print(f\"Saved CSVs for lore {lore_id} / root {root_id}\")"
   ],
   "id": "4q-sNfDf-uqr"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IHvPRBs7-uqs"
   },
   "source": [
    "---\n",
    "## 8 \u00b7 Listing Proofread Segments\n",
    "\n",
    "The function below returns all segments that have received at least one proofreading edit (merge or split), with edit counts and soma coordinates."
   ],
   "id": "IHvPRBs7-uqs"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "G8G5sCen-uqs"
   },
   "outputs": [],
   "source": [
    "#@title Define helper: get_proofread_segments\n",
    "from datetime import datetime, timezone\n",
    "\n",
    "def get_proofread_segments(client, since=None, limit=None):\n",
    "    \"\"\"\n",
    "    Return a DataFrame of all segments with at least one edit.\n",
    "\n",
    "    Args:\n",
    "        since: datetime \u2014 only include edits after this date (default: all)\n",
    "        limit: int \u2014 cap the number of edited segments to fetch (default: all)\n",
    "\n",
    "    Columns: root_id, n_merges, n_splits, n_total_edits,\n",
    "             first_edit, last_edit, x, y, z  (soma voxel coordinate)\n",
    "    \"\"\"\n",
    "    ts_start = since or datetime(2020, 1, 1, tzinfo=timezone.utc)\n",
    "    ts_end   = datetime.now(tz=timezone.utc)\n",
    "\n",
    "    _, new_roots = client.chunkedgraph.get_delta_roots(ts_start, ts_end)\n",
    "    edited_roots = np.unique(new_roots).tolist()\n",
    "    print(f\"Found {len(edited_roots)} edited segments total\")\n",
    "\n",
    "    if limit:\n",
    "        edited_roots = edited_roots[:limit]\n",
    "        print(f\"Fetching first {limit}\")\n",
    "\n",
    "    BATCH, all_logs = 200, {}\n",
    "    for i in range(0, len(edited_roots), BATCH):\n",
    "        batch = edited_roots[i : i + BATCH]\n",
    "        all_logs.update(\n",
    "            client.chunkedgraph.get_tabular_change_log(batch, filtered=False)\n",
    "        )\n",
    "        print(f\"  fetched {i + len(batch)} / {len(edited_roots)}\")\n",
    "\n",
    "    frames = []\n",
    "    for rid, df in all_logs.items():\n",
    "        if df.empty:\n",
    "            continue\n",
    "        df[\"timestamp\"] = pd.to_datetime(df[\"timestamp\"], utc=True)\n",
    "        frames.append(pd.DataFrame({\n",
    "            \"root_id\":        rid,\n",
    "            \"n_merges\":       int(df[\"is_merge\"].sum()),\n",
    "            \"n_splits\":       int((~df[\"is_merge\"]).sum()),\n",
    "            \"n_total_edits\":  len(df),\n",
    "            \"first_edit\":     df[\"timestamp\"].min(),\n",
    "            \"last_edit\":      df[\"timestamp\"].max(),\n",
    "        }, index=[0]))\n",
    "\n",
    "    if not frames:\n",
    "        return pd.DataFrame()\n",
    "\n",
    "    edit_summary = pd.concat(frames, ignore_index=True)\n",
    "\n",
    "    soma_locs = client.materialize.query_table(\n",
    "        \"somas\",\n",
    "        filter_in_dict={\"pt_root_id\": edit_summary[\"root_id\"].tolist()},\n",
    "        select_columns=[\"pt_root_id\", \"pt_position\"],\n",
    "    ).rename(columns={\"pt_root_id\": \"root_id\"})\n",
    "    soma_locs[[\"x\", \"y\", \"z\"]] = pd.DataFrame(\n",
    "        soma_locs[\"pt_position\"].tolist(), index=soma_locs.index\n",
    "    )\n",
    "    soma_locs.drop(columns=[\"pt_position\"], inplace=True)\n",
    "\n",
    "    return (\n",
    "        edit_summary.merge(soma_locs, on=\"root_id\", how=\"left\")\n",
    "        .sort_values(\"n_total_edits\", ascending=False)\n",
    "        .reset_index(drop=True)\n",
    "    )"
   ],
   "id": "G8G5sCen-uqs"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "glHNEnTU-uqs",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "cc9d9c32-17e9-473e-ab35-8ae1338384aa"
   },
   "outputs": [],
   "source": [
    "#@title Run: list all proofread segments\n",
    "df = get_proofread_segments(client, limit = 500)\n",
    "\n",
    "print(\"=\" * 50)\n",
    "print(\"PROOFREADING SUMMARY\")\n",
    "print(\"=\" * 50)\n",
    "print(f\"  Total edited segments : {len(df)}\")\n",
    "print(f\"  Total merges          : {df.n_merges.sum()}\")\n",
    "print(f\"  Total splits          : {df.n_splits.sum()}\")\n",
    "print(f\"  Total operations      : {df.n_total_edits.sum()}\")\n",
    "print(f\"  Segments with >10 edits: {(df.n_total_edits > 10).sum()}\")\n",
    "if not df.empty:\n",
    "    print(f\"  First edit            : {df.first_edit.min()}\")\n",
    "    print(f\"  Latest edit           : {df.last_edit.max()}\")\n",
    "print(\"=\" * 50)\n",
    "print()\n",
    "print(df.head(20))\n",
    "\n",
    "# Uncomment to save:\n",
    "# df.to_csv(\"proofread_segments.csv\", index=False)"
   ],
   "id": "glHNEnTU-uqs"
  }
 ]
}