{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a9be3e89",
   "metadata": {},
   "source": [
    "# GraphDataset / SubgraphDataset with bibliographic data (OpenAlex)\n",
    "\n",
    "This notebook shows `cvcdocdb`'s `torch_dataloader` on top of the sample\n",
    "bibliographic dataset (`cvcdocdb.exemples.load_bibliografia_openalex`):\n",
    "\n",
    "- `GraphDataset` + `GraphDataLoader`: flat streaming over nodes (`Paper`, `Author`).\n",
    "- `SubgraphDataset` + `PyGDataLoader`: k-hop ego subgraphs as\n",
    "  `torch_geometric.data.Data` objects, with an embedding per `Paper` (`data.x` and `data.emb`).\n",
    "- `MetaPath2Vec`: quick embedding training on the heterogeneous\n",
    "  `Author`-`Paper` graph, following the metapath `author -writes-> paper -cites-> paper -written_by-> author`.\n",
    "\n",
    "Since the OpenAlex dataset doesn't carry real embeddings, we first generate a\n",
    "deterministic embedding (hash of the title) purely for demonstration — and\n",
    "later train a real embedding with MetaPath2Vec, which you can use to replace it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f67e1098",
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib.util\n",
    "\n",
    "# Comprovació de presència del paquet\n",
    "package_to_check = 'cvcdocdb'\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 cvcdocdb\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ó.')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24651550",
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib.util\n",
    "\n",
    "for package_to_check in ('torch', 'torch_geometric'):\n",
    "    spec = importlib.util.find_spec(package_to_check)\n",
    "    if spec is None:\n",
    "        print(f'⚠️ {package_to_check} no està instal·lat. Iniciant instal·lació...')\n",
    "        %pip install -q {package_to_check}\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ó.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d455a338",
   "metadata": {},
   "source": [
    "## 1. Load the bibliographic dataset into `NetworkXGraph`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d298b6df",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cvcdocdb import NetworkXGraph\n",
    "from cvcdocdb.exemples import load_bibliografia_openalex\n",
    "\n",
    "store = NetworkXGraph()\n",
    "stats = load_bibliografia_openalex(store, query=\"graph database\", per_page=15)\n",
    "print(\"Dataset stats:\", stats)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8a6a4ed",
   "metadata": {},
   "source": [
    "## 2. Add a demo embedding to each `Paper`\n",
    "\n",
    "`load_bibliografia_openalex` doesn't generate embeddings. Here we compute a\n",
    "deterministic one (hash of the title -> normalised vector) and add it as the\n",
    "`embedding` attribute of each `Paper` node with `insertNode(..., update=True)`,\n",
    "which merges the attribute without touching the rest of the node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8b78203",
   "metadata": {},
   "outputs": [],
   "source": [
    "import hashlib\n",
    "\n",
    "from cvcdocdb.base import Node\n",
    "\n",
    "EMB_DIM = 16\n",
    "\n",
    "\n",
    "def toy_embedding(text: str, dim: int = EMB_DIM) -> list[float]:\n",
    "    \"\"\"Deterministic embedding (demo only) derived from the hash of the text.\"\"\"\n",
    "    digest = hashlib.sha256(text.encode(\"utf-8\")).digest()\n",
    "    raw = [digest[i % len(digest)] / 255.0 for i in range(dim)]\n",
    "    norm = sum(v * v for v in raw) ** 0.5 or 1.0\n",
    "    return [v / norm for v in raw]\n",
    "\n",
    "\n",
    "paper_ids = store.find_nodes_by_property(\"main_label\", \"Paper\")\n",
    "for nid in paper_ids:\n",
    "    attrs = store.get_node_attrs(nid)\n",
    "    title = attrs.get(\"title\", \"\")\n",
    "    openalex_id = attrs.get(\"openalex_id\")\n",
    "    embedding = toy_embedding(title)\n",
    "    store.insertNode(\n",
    "        Node(pk={\"openalex_id\": openalex_id}, main_label=\"Paper\", embedding=embedding),\n",
    "        update=True,\n",
    "    )\n",
    "\n",
    "print(f\"Added embeddings to {len(paper_ids)} papers (dim={EMB_DIM}).\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1429c5f5",
   "metadata": {},
   "source": [
    "## 3. Flat streaming with `GraphDataset` / `GraphDataLoader`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46581d27",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cvcdocdb.torch_dataloader import GraphDataset, GraphDataLoader\n",
    "\n",
    "flat_ds = GraphDataset(store, label_filter=\"Paper\")\n",
    "flat_loader = GraphDataLoader(flat_ds, batch_size=8, num_workers=0)\n",
    "\n",
    "batch = next(iter(flat_loader))\n",
    "print(\"node_id:\", batch[\"node_id\"])\n",
    "print(\"main_label:\", batch[\"main_label\"][:3], \"...\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "854da3f2",
   "metadata": {},
   "source": [
    "## 4. Ego subgraphs with embeddings via `SubgraphDataset` / `PyGDataLoader`\n",
    "\n",
    "Each `Paper` is a seed; the k-hop subgraph includes its co-authors and\n",
    "cited/citing papers. `vector_attr=\"embedding\"` merges the embedding into\n",
    "`data.x` and also keeps it accessible separately on `data.emb`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c6d4cd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cvcdocdb.torch_dataloader import SubgraphDataset, PyGDataLoader\n",
    "\n",
    "sub_ds = SubgraphDataset(\n",
    "    store,\n",
    "    hops=1,\n",
    "    node_attrs=[\"year\"],\n",
    "    vector_attr=\"embedding\",\n",
    "    label_filter=\"Paper\",\n",
    "    missing=0.0,\n",
    ")\n",
    "sub_loader = PyGDataLoader(sub_ds, batch_size=4)\n",
    "\n",
    "pyg_batch = next(iter(sub_loader))\n",
    "print(\"x:\", pyg_batch.x.shape)          # [N, 1 (year) + 16 (embedding)]\n",
    "print(\"emb:\", pyg_batch.emb.shape)      # [N, 16]\n",
    "print(\"edge_index:\", pyg_batch.edge_index.shape)\n",
    "print(\"node_ids:\", pyg_batch.node_ids[:10])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aea98ac",
   "metadata": {},
   "source": [
    "## 5. Training embeddings with `MetaPath2Vec`\n",
    "\n",
    "Instead of section 2's \"toy\" embedding, we train a real embedding with\n",
    "`torch_geometric.nn.MetaPath2Vec` on the full heterogeneous `Author`/`Paper`\n",
    "graph, following the metapath:\n",
    "\n",
    "```\n",
    "Author --AUTHORED--> Paper --CITES--> Paper --AUTHORED_BY--> Author\n",
    "```\n",
    "\n",
    "`to_hetero_edge_index_dict` loads the whole graph in a single pass over the\n",
    "edges (no node attributes, as lightweight as possible) and returns the\n",
    "`edge_index_dict`, the node count per type, and the original-id -> local-index\n",
    "mapping that `MetaPath2Vec` needs.\n",
    "\n",
    "Since the dataset is small (an OpenAlex sample of ~15 papers), training is\n",
    "brief (few epochs) and purely illustrative — on a real dataset you should\n",
    "increase `walk_length`, `walks_per_node`, `context_size`, and the number of\n",
    "epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24d5d9eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch_geometric.nn import MetaPath2Vec\n",
    "\n",
    "from cvcdocdb.torch_dataloader import to_hetero_edge_index_dict\n",
    "\n",
    "edge_index_dict, num_nodes_dict, node_maps = to_hetero_edge_index_dict(store)\n",
    "print(\"Node types:\", num_nodes_dict)\n",
    "print(\"Relation types:\", list(edge_index_dict))\n",
    "\n",
    "# Add the reverse Paper->Author relation so the metapath can close the cycle.\n",
    "authored = edge_index_dict[(\"Author\", \"AUTHORED\", \"Paper\")]\n",
    "edge_index_dict[(\"Paper\", \"AUTHORED_BY\", \"Author\")] = authored.flip(0)\n",
    "\n",
    "if (\"Paper\", \"CITES\", \"Paper\") not in edge_index_dict:\n",
    "    # Small sample: there may be no intra-sample citations. Add self-loops\n",
    "    # so the paper->paper step of the metapath stays traversable.\n",
    "    n_papers = num_nodes_dict[\"Paper\"]\n",
    "    idx = torch.arange(n_papers, dtype=torch.long)\n",
    "    edge_index_dict[(\"Paper\", \"CITES\", \"Paper\")] = torch.stack([idx, idx])\n",
    "\n",
    "metapath = [\n",
    "    (\"Author\", \"AUTHORED\", \"Paper\"),\n",
    "    (\"Paper\", \"CITES\", \"Paper\"),\n",
    "    (\"Paper\", \"AUTHORED_BY\", \"Author\"),\n",
    "]\n",
    "\n",
    "model = MetaPath2Vec(\n",
    "    edge_index_dict,\n",
    "    embedding_dim=16,\n",
    "    metapath=metapath,\n",
    "    walk_length=4,\n",
    "    context_size=2,\n",
    "    walks_per_node=3,\n",
    "    num_negative_samples=3,\n",
    "    num_nodes_dict=num_nodes_dict,\n",
    "    sparse=True,\n",
    ")\n",
    "\n",
    "loader = model.loader(batch_size=8, shuffle=True)\n",
    "optimizer = torch.optim.SparseAdam(list(model.parameters()), lr=0.01)\n",
    "\n",
    "model.train()\n",
    "for epoch in range(5):\n",
    "    total_loss = 0.0\n",
    "    for pos_rw, neg_rw in loader:\n",
    "        optimizer.zero_grad()\n",
    "        loss = model.loss(pos_rw, neg_rw)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += float(loss)\n",
    "    print(f\"epoch {epoch + 1}: loss={total_loss / max(len(loader), 1):.4f}\")\n",
    "\n",
    "paper_embeddings = model(\"Paper\").detach()\n",
    "print(\"Paper embeddings (MetaPath2Vec):\", paper_embeddings.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38787e26",
   "metadata": {},
   "source": [
    "With the learned embedding we can, for example, look at which sample\n",
    "papers are most similar to each other (cosine similarity) — a typical use of\n",
    "these embeddings is to feed them back into the graph\n",
    "(`insertNode(..., update=True)`) to use them with `SubgraphDataset(vector_attr=...)`\n",
    "as in section 4."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6f5ec71",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "sample = paper_embeddings[: min(5, len(paper_ids))]\n",
    "sim = F.cosine_similarity(sample.unsqueeze(1), sample.unsqueeze(0), dim=-1)\n",
    "print(\"Cosine similarity between the first papers:\")\n",
    "print(sim)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cee4abe",
   "metadata": {},
   "source": [
    "## 6. Link prediction: predicting citations between papers\n",
    "\n",
    "With `Paper` embeddings already trained, let's run a simple **link\n",
    "prediction** example: predicting whether a `CITES` relation exists between\n",
    "two papers. Since `MetaPath2Vec` only learns *node* embeddings, we derive an\n",
    "edge representation by combining both endpoints with\n",
    "`edge_embeddings(..., op=\"dot\")` — the dot product is directly a\n",
    "similarity/existence score.\n",
    "\n",
    "We split the `CITES` edges by year (`split_edges_by_node_property`): training\n",
    "on \"older\" paper pairs and validating on \"recent\" ones. Since we also need\n",
    "negative examples (pairs with no citation), we use\n",
    "`torch_geometric.utils.negative_sampling`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69244792",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch_geometric.utils import negative_sampling\n",
    "\n",
    "from cvcdocdb.torch_dataloader import edge_embeddings, split_edges_by_node_property\n",
    "\n",
    "cites_edge_index = edge_index_dict[(\"Paper\", \"CITES\", \"Paper\")]\n",
    "\n",
    "# Publication year per paper, aligned with node_maps[\"Paper\"] (local index).\n",
    "local_to_orig_paper = {v: k for k, v in node_maps[\"Paper\"].items()}\n",
    "years = torch.tensor([\n",
    "    float((store.get_node_attrs(local_to_orig_paper[i]) or {}).get(\"year\") or 0)\n",
    "    for i in range(num_nodes_dict[\"Paper\"])\n",
    "])\n",
    "\n",
    "cutoff = years.median()\n",
    "splits = split_edges_by_node_property(\n",
    "    cites_edge_index, years,\n",
    "    train=lambda s, d: torch.maximum(s, d) <= cutoff,\n",
    "    test=lambda s, d: torch.minimum(s, d) >= cutoff,\n",
    ")\n",
    "print(\"Cutoff (median year):\", cutoff.item())\n",
    "print(\"CITES edges — train:\", splits[\"train\"].shape[1], \" test:\", splits[\"test\"].shape[1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "594ed3c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_papers = num_nodes_dict[\"Paper\"]\n",
    "\n",
    "\n",
    "def score_split(pos_edge_index, name):\n",
    "    if pos_edge_index.shape[1] == 0:\n",
    "        print(f\"{name}: no positive edges, skipping.\")\n",
    "        return\n",
    "    neg_edge_index = negative_sampling(\n",
    "        edge_index=cites_edge_index,\n",
    "        num_nodes=n_papers,\n",
    "        num_neg_samples=pos_edge_index.shape[1],\n",
    "    )\n",
    "    pos_scores = edge_embeddings(paper_embeddings, pos_edge_index, op=\"dot\")\n",
    "    neg_scores = edge_embeddings(paper_embeddings, neg_edge_index, op=\"dot\")\n",
    "\n",
    "    preds = torch.cat([pos_scores, neg_scores])\n",
    "    labels = torch.cat([torch.ones_like(pos_scores), torch.zeros_like(neg_scores)])\n",
    "    accuracy = ((preds > 0).float() == labels).float().mean()\n",
    "\n",
    "    print(f\"{name}: mean positive score={pos_scores.mean():.3f}  \"\n",
    "          f\"negative={neg_scores.mean():.3f}  accuracy(threshold=0)={accuracy:.2%}\")\n",
    "\n",
    "\n",
    "score_split(splits[\"train\"], \"train\")\n",
    "score_split(splits[\"test\"], \"test\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "880bb77d",
   "metadata": {},
   "outputs": [],
   "source": [
    "store.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (cvcdocdb)",
   "language": "python",
   "name": "cvcdocdb"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}