{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vector Search with HNSW\n",
    "\n",
    "This notebook demonstrates the vector indexing and nearest-neighbor search capabilities\n",
    "of `NetworkXGraph`. It uses [hnswlib](https://github.com/nmslib/hnswlib) for approximate\n",
    "nearest-neighbor (ANN) queries.\n",
    "\n",
    "Key methods covered:\n",
    "- `enable_vector_index()` — create an ANN index on a node property\n",
    "- `query_vector_index()` — find the k nearest nodes to a query vector\n",
    "- Supported distance spaces: `cosine`, `l2`, `ip` (inner product)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib.util\n",
    "\n",
    "package_to_check = 'drm'\n",
    "spec = importlib.util.find_spec(package_to_check)\n",
    "\n",
    "if spec is None:\n",
    "    print(f'⚠️ {package_to_check} no està instal·lat. Iniciant instal·lació...')\n",
    "    %pip install -q --upgrade drm-tools\n",
    "    print(\"✅ Instal·lació completada. L'estat del kernel PODRIA requerir un reinici.\")\n",
    "else:\n",
    "    print(f'✅ {package_to_check} ja està present al sistema. Saltant instal·lació.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "try:\n",
    "    import hnswlib\n",
    "    print(\"✅ hnswlib is available\")\n",
    "except ImportError:\n",
    "    print(\"⚠️ hnswlib not installed. Run: pip install hnswlib\")\n",
    "    hnswlib = None\n",
    "\n",
    "from drm import NetworkXGraph, Node\n",
    "\n",
    "graph = NetworkXGraph()\n",
    "\n",
    "# Create a small set of items with 4-dimensional feature vectors\n",
    "# Each item has a 'features' property that will be indexed for vector search\n",
    "items = [\n",
    "    {\"name\": \"apple\",    \"features\": [0.9, 0.1, 0.05, 0.0],   \"category\": \"fruit\"},\n",
    "    {\"name\": \"banana\",   \"features\": [0.85, 0.15, 0.0, 0.0],  \"category\": \"fruit\"},\n",
    "    {\"name\": \"carrot\",   \"features\": [0.0, 0.9, 0.05, 0.05],  \"category\": \"vegetable\"},\n",
    "    {\"name\": \"broccoli\", \"features\": [0.05, 0.85, 0.05, 0.05], \"category\": \"vegetable\"},\n",
    "    {\"name\": \"salmon\",   \"features\": [0.1, 0.1, 0.7, 0.1],    \"category\": \"protein\"},\n",
    "    {\"name\": \"chicken\",  \"features\": [0.1, 0.1, 0.75, 0.05],  \"category\": \"protein\"},\n",
    "    {\"name\": \"rice\",     \"features\": [0.2, 0.2, 0.1, 0.5],    \"category\": \"grain\"},\n",
    "]\n",
    "\n",
    "node_ids = {}\n",
    "for item in items:\n",
    "    node = Node(pk={\"name\": item[\"name\"]}, main_label=\"Food\", features=item[\"features\"], category=item[\"category\"])\n",
    "    nid = graph.insertNode(node)\n",
    "    node_ids[item[\"name\"]] = nid\n",
    "\n",
    "print(f\"Inserted {len(items)} food items\")\n",
    "print(\"Node ids:\", node_ids)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Enabling a vector index\n",
    "\n",
    "Call `enable_vector_index(property_name, dimensions, space)` to create an HNSW index\n",
    "on a specific node property. The property values must be 1D vectors of the given\n",
    "dimensionality.\n",
    "\n",
    "Supported `space` values:\n",
    "- `\"cosine\"` — cosine similarity (default, recommended for normalized vectors)\n",
    "- `\"l2\"` — Euclidean distance\n",
    "- `\"ip\"` — inner product"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph.enable_vector_index(\n",
    "    property_name=\"features\",\n",
    "    dimensions=4,\n",
    "    space=\"cosine\",\n",
    ")\n",
    "print(\"Vector index enabled for property 'features' (4D, cosine distance)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying nearest neighbors\n",
    "\n",
    "`query_vector_index(property_name, vector, top_k)` returns a list of\n",
    "`(node_id, distance)` tuples sorted by ascending distance.\n",
    "\n",
    "We'll query for the nearest neighbors of a \"mystery\" food item."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lookup_name(node_id):\n",
    "    \"\"\"Reverse lookup: find the food name for a node id.\"\"\"\n",
    "    for name, nid in node_ids.items():\n",
    "        if nid == node_id:\n",
    "            return name\n",
    "    return f\"id={node_id}\"\n",
    "\n",
    "def show_results(query_vec, top_k=3):\n",
    "    results = graph.query_vector_index(\n",
    "        property_name=\"features\",\n",
    "        vector=query_vec,\n",
    "        top_k=top_k,\n",
    "    )\n",
    "    print(f\"Query: {query_vec}\")\n",
    "    print(f\"Top {top_k} nearest neighbors:\")\n",
    "    for nid, dist in results:\n",
    "        name = lookup_name(nid)\n",
    "        print(f\"  {name:10s} distance={dist:.4f}\")\n",
    "    return results\n",
    "\n",
    "# Query 1: Find fruits (similar to apple)\n",
    "print(\"=== Query 1: Find fruits ===\")\n",
    "show_results([0.9, 0.05, 0.0, 0.0], top_k=3)\n",
    "\n",
    "print()\n",
    "\n",
    "# Query 2: Find vegetables\n",
    "print(\"=== Query 2: Find vegetables ===\")\n",
    "show_results([0.0, 0.8, 0.1, 0.1], top_k=3)\n",
    "\n",
    "print()\n",
    "\n",
    "# Query 3: Find protein\n",
    "print(\"=== Query 3: Find protein ===\")\n",
    "show_results([0.05, 0.05, 0.8, 0.1], top_k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding new nodes updates the index automatically\n",
    "\n",
    "When you insert a new node with the indexed property, it is automatically\n",
    "added to the HNSW index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a new food item\n",
    "mango = Node(pk={\"name\": \"mango\"}, main_label=\"Food\", features=[0.88, 0.1, 0.02, 0.0], category=\"fruit\")\n",
    "graph.insertNode(mango)\n",
    "node_ids[\"mango\"] = graph.checkNode(mango)\n",
    "print(\"Added 'mango' to the graph\")\n",
    "\n",
    "# Query again — mango should appear near apple and banana\n",
    "print()\n",
    "print(\"=== Query after adding mango ===\")\n",
    "show_results([0.9, 0.1, 0.0, 0.0], top_k=4)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple vector indexes\n",
    "\n",
    "You can create separate indexes on different properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a second index on a different property (if you had 2D embedding)\n",
    "# For this demo, we'll just show the API:\n",
    "# graph.enable_vector_index(\n",
    "#     property_name=\"other_features\",\n",
    "#     dimensions=8,\n",
    "#     space=\"l2\",\n",
    "# )\n",
    "print(\"Multiple vector indexes are supported. See the API for details.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cleanup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph.close()\n",
    "print(\"Graph closed.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}