{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przewidywanie Interakcji Leków z wykorzystaniem Graph Neural Networks\n",
    "\n",
    "## Wprowadzenie\n",
    "W tym notatniku przedstawimy implementację modelu do przewidywania interakcji między lekami (Drug-Drug Interactions, DDI) wykorzystując zbiór danych z biblioteki OGB (Open Graph Benchmark).\n",
    "\n",
    "### Cel\n",
    "Głównym celem jest zbudowanie modelu, który potrafi przewidzieć, czy między dwoma lekami występuje interakcja. Jest to przykład zadania przewidywania połączeń w grafie (link prediction).\n",
    "\n",
    "## Struktura implementacji:\n",
    "1. Wczytanie i przygotowanie danych z OGB\n",
    "2. Przygotowanie próbek negatywnych\n",
    "3. Implementacja modelu GNN do generowania osadzeń (ang. *emgedding*)\n",
    "4. Implementacja predyktora połączeń\n",
    "5. Implementacja modułu treningowego PyTorch Lightning\n",
    "6. Trening i ewaluacja modelu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/filip/projects/grafowe_sieci_neuronowe/.venv/lib/python3.12/site-packages/outdated/__init__.py:36: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n",
      "  from pkg_resources import parse_version\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "import joblib\n",
    "import torch.nn as nn\n",
    "import networkx as nx\n",
    "import optuna\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.data import Data\n",
    "import torch_geometric as pyg\n",
    "from torch_geometric.transforms import RandomLinkSplit\n",
    "from ogb.linkproppred import PygLinkPropPredDataset\n",
    "import pytorch_lightning as pl\n",
    "import torchmetrics as tm\n",
    "from torch_geometric.nn import GCNConv, SAGEConv, GATConv\n",
    "import numpy as np\n",
    "import json\n",
    "from sklearn.metrics import classification_report"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przygotowanie Danych\n",
    "\n",
    "Wykorzystamy zbiór danych ogbl-ddi z biblioteki OGB, który zawiera informacje o interakcjach między lekami. \n",
    "Każdy węzeł reprezentuje lek, a krawędzie oznaczają udokumentowane interakcje między lekami."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://snap.stanford.edu/ogb/data/linkproppred/ddi.zip\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloaded 0.04 GB: 100%|██████████| 46/46 [00:18<00:00,  2.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting dataset/ddi.zip\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading necessary files...\n",
      "This might take a while.\n",
      "Processing graphs...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00, 83.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converting graphs into PyG objects...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00, 8439.24it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "# Poniższe importy i rejestracja klas są wymagane dla kompatybilności zbioru danych OGB z Pytorchem.\n",
    "# Szczegóły: https://github.com/snap-stanford/ogb/issues/497\n",
    "# W chwili przygotowania niniejszego materiały OGB nie jest kompatybilne z nowszymi wersjami PyTorcha - 2.6+\n",
    "from torch_geometric.data.storage import GlobalStorage\n",
    "from torch_geometric.data.data import DataEdgeAttr, DataTensorAttr\n",
    "torch.serialization.add_safe_globals([GlobalStorage, DataEdgeAttr, DataTensorAttr])\n",
    "\n",
    "# Wczytanie zbioru danych\n",
    "dataset = PygLinkPropPredDataset(name=\"ogbl-ddi\")\n",
    "original_data = dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Próbka - spójny pograf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ze względu na znaczną wielkość i gęstość oryginalnego grafu, jego pełne przetwarzanie wymagałoby \n",
    "istotnych zasobów obliczeniowych. Na potrzeby niniejszego przewodnika ograniczymy się do \n",
    "reprezentatywnej próbki danych - wykorzystamy 20% wierzchołków z największej spójnej składowej grafu.\n",
    "Pozwoli to na efektywną demonstrację koncepcji i technik przy zachowaniu rozsądnych wymagań\n",
    "sprzętowych."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "`recreate=True` powinno być ustawione, gdy uruchamiamy kod po raz pierwszy, aby odtworzyć proces przetwarzania.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "recreate = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "Uwaga, poniższe operacje mogą zająć więcej czasu.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Liczba wierzchołków (20% oryginalnego zbioru): 845\n",
      "Liczba krawędzi w podgrafie: 170476\n",
      "Liczba cech wierzchołków: 6\n"
     ]
    }
   ],
   "source": [
    "if recreate:\n",
    "    # Konwersja do formatu NetworkX\n",
    "    data_nx = pyg.utils.to_networkx(original_data)\n",
    "\n",
    "    # Znajdujemy największą spójną składową i tworzymy z niej mniejszy podgraf\n",
    "    connected_components = list(nx.connected_components(data_nx.to_undirected()))\n",
    "    largest_cc = max(connected_components, key=len)\n",
    "\n",
    "    # Wybieramy losowy podzbiór wierzchołków z największej spójnej składowej (20% oryginalnego rozmiaru)\n",
    "    subset_size = int(len(largest_cc) * 0.2)\n",
    "    subset_nodes = list(largest_cc)[:subset_size]\n",
    "\n",
    "    # Tworzymy podgraf z wybranych wierzchołków\n",
    "    subgraph = data_nx.subgraph(subset_nodes)\n",
    "\n",
    "    # Upewniamy się, że podgraf jest spójny biorąc największą spójną składową\n",
    "    connected_subgraph_nodes = max(\n",
    "        nx.connected_components(subgraph.to_undirected()), key=len\n",
    "    )\n",
    "    connected_subgraph = nx.subgraph(data_nx, connected_subgraph_nodes).to_undirected()\n",
    "\n",
    "    assert nx.is_connected(connected_subgraph.to_undirected())\n",
    "\n",
    "    data = pyg.utils.from_networkx(connected_subgraph)\n",
    "\n",
    "    # Obliczamy stopnie wierzchołków\n",
    "    degrees = pyg.utils.degree(data.edge_index[0], data.num_nodes)\n",
    "\n",
    "    # Tworzymy graf networkx do obliczenia miar centralności\n",
    "    betweenness = torch.tensor(\n",
    "        list(nx.betweenness_centrality(connected_subgraph).values())\n",
    "    )\n",
    "    closeness = torch.tensor(list(nx.closeness_centrality(connected_subgraph).values()))\n",
    "    pagerank = torch.tensor(list(nx.pagerank(connected_subgraph, max_iter=50).values()))\n",
    "    clustering_coeff = torch.tensor(list(nx.clustering(connected_subgraph).values()))\n",
    "    ones = torch.ones(data.num_nodes)\n",
    "\n",
    "    # Łączymy wszystkie cechy w jedną macierz\n",
    "    data.x = torch.stack(\n",
    "        [degrees, betweenness, closeness, pagerank, clustering_coeff, ones], dim=1\n",
    "    ).float()\n",
    "\n",
    "    # Wyświetlenie informacji o podgrafie\n",
    "    print(f\"Liczba wierzchołków (20% oryginalnego zbioru): {data.num_nodes}\")\n",
    "    print(f\"Liczba krawędzi w podgrafie: {data.edge_index.size(1)}\")\n",
    "    print(\n",
    "        f\"Liczba cech wierzchołków: {data.num_features if hasattr(data, 'num_features') else 'brak cech'}\"\n",
    "    )\n",
    "\n",
    "    # Podział na zbiór treningowy, walidacyjny i testowy\n",
    "    pyg.seed_everything(123)\n",
    "    transform = RandomLinkSplit(\n",
    "        num_val=0.1,\n",
    "        num_test=0.1,\n",
    "        neg_sampling_ratio=1.0,  # Stosunek próbek negatywnych do pozytywnych\n",
    "        add_negative_train_samples=True,\n",
    "        is_undirected=True,\n",
    "    )\n",
    "\n",
    "    train_data, val_data, test_data = transform(data)\n",
    "    # Zapisujemy dane do plików\n",
    "    torch.save(data, \"data.pt\")\n",
    "    torch.save(train_data, \"train_data.pt\")\n",
    "    torch.save(val_data, \"val_data.pt\")\n",
    "    torch.save(test_data, \"test_data.pt\")\n",
    "else:\n",
    "    data = torch.load(\"data.pt\")\n",
    "    train_data = torch.load(\"train_data.pt\", weights_only=False)\n",
    "    val_data = torch.load(\"val_data.pt\", weights_only=False)\n",
    "    test_data = torch.load(\"test_data.pt\", weights_only=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n",
      "False\n",
      "False\n"
     ]
    }
   ],
   "source": [
    "# Sprawdzenie, czy wśród negatywnych próbek są krawędzie rzeczywiste (nie powinno być - false positives)\n",
    "set_of_all_real_edges = set([tuple(e) for e in data.edge_index.t().tolist()])\n",
    "\n",
    "neg_train_edges = train_data.edge_label_index[:, train_data.edge_label == 0]\n",
    "neg_train_edges = set([tuple(e) for e in neg_train_edges.t().tolist()])\n",
    "\n",
    "neg_val_edges = val_data.edge_label_index[:, val_data.edge_label == 0]\n",
    "neg_val_edges = set([tuple(e) for e in neg_val_edges.t().tolist()])\n",
    "\n",
    "neg_test_edges = test_data.edge_label_index[:, test_data.edge_label == 0]\n",
    "neg_test_edges = set([tuple(e) for e in neg_test_edges.t().tolist()])\n",
    "\n",
    "print(any([ne in set_of_all_real_edges for ne in neg_train_edges]))\n",
    "print(any([ne in set_of_all_real_edges for ne in neg_val_edges]))\n",
    "print(any([ne in set_of_all_real_edges for ne in neg_test_edges]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(edge_index=[2, 136384], num_nodes=845, x=[845, 6], edge_label=[136384], edge_label_index=[2, 136384])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(edge_index=[2, 136384], num_nodes=845, x=[845, 6], edge_label=[17046], edge_label_index=[2, 17046])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(edge_index=[2, 153430], num_nodes=845, x=[845, 6], edge_label=[17046], edge_label_index=[2, 17046])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = pyg.loader.LinkNeighborLoader(\n",
    "    train_data,\n",
    "    batch_size=1024,\n",
    "    num_neighbors=[-1, 15, 10, 5],\n",
    "    shuffle=True, \n",
    "    edge_label_index=train_data.edge_label_index,\n",
    "    edge_label=train_data.edge_label,\n",
    "    num_workers=8)\n",
    "\n",
    "val_loader = pyg.loader.DataLoader(\n",
    "    [val_data],\n",
    "    batch_size=1)\n",
    "\n",
    "test_loader = pyg.loader.DataLoader(\n",
    "    [test_data], batch_size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "train_batch_props = []\n",
    "for b in train_loader:\n",
    "    train_batch_props.append(pd.Series(b.edge_label).value_counts(normalize=True))\n",
    "    \n",
    "val_batch_props = []\n",
    "for b in val_loader:\n",
    "    val_batch_props.append(pd.Series(b.edge_label).value_counts(normalize=True))\n",
    "    \n",
    "test_batch_props = []\n",
    "for b in test_loader:\n",
    "    test_batch_props.append(pd.Series(b.edge_label).value_counts(normalize=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_batch_props = pd.concat(train_batch_props, axis=1).mean(axis=1)\n",
    "val_batch_props = pd.concat(val_batch_props, axis=1).mean(axis=1)\n",
    "test_batch_props = pd.concat(test_batch_props, axis=1).mean(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Średnie proporcje klas w batchach treningowych: [0.5, 0.5] \n",
      "\n",
      "\n",
      "Średnie proporcje klas w batchach walidacyjnych: [0.5, 0.5] \n",
      "\n",
      "\n",
      "Średnie proporcje klas w batchach testowych: [0.5, 0.5] \n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(f\"Średnie proporcje klas w batchach treningowych: {train_batch_props.values.round(2).tolist()} \\n\\n\")\n",
    "print(f\"Średnie proporcje klas w batchach walidacyjnych: {val_batch_props.values.round(2).tolist()} \\n\\n\")\n",
    "print(f\"Średnie proporcje klas w batchach testowych: {test_batch_props.values.round(2).tolist()} \\n\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modele"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predyktor Połączeń\n",
    "\n",
    "Implementujemy model do przewidywania istnienia połączeń między parami wierzchołków na podstawie ich embeddingów."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinkPredictor(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Model do przewidywania połączeń między wierzchołkami.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    in_channels : int\n",
    "        Wymiar embeddingów wejściowych\n",
    "    hidden_channels : int\n",
    "        Wymiar warstw ukrytych\n",
    "    num_layers : int\n",
    "        Liczba warstw MLP\n",
    "    dropout : float\n",
    "        Współczynnik dropout\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, hidden_channels, num_layers, dropout):\n",
    "        super().__init__()\n",
    "\n",
    "        self.lins = torch.nn.ModuleList()\n",
    "        self.lins.append(torch.nn.Linear(in_channels * 2, hidden_channels))\n",
    "        for _ in range(num_layers - 2):\n",
    "            self.lins.append(torch.nn.Linear(hidden_channels, hidden_channels))\n",
    "        self.lins.append(torch.nn.Linear(hidden_channels, 1))\n",
    "\n",
    "        self.dropout = dropout\n",
    "\n",
    "    def forward(self, x_i, x_j):\n",
    "        \"\"\"\n",
    "        Propagacja w przód.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x_i : torch.Tensor\n",
    "            Embeddingi pierwszych wierzchołków w parach\n",
    "        x_j : torch.Tensor\n",
    "            Embeddingi drugich wierzchołków w parach\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Prawdopodobieństwa istnienia połączeń\n",
    "        \"\"\"\n",
    "        x = torch.cat([x_i, x_j], dim=-1)\n",
    "\n",
    "        for lin in self.lins[:-1]:\n",
    "            x = lin(x)\n",
    "            x = F.relu(x)\n",
    "            x = F.dropout(x, p=self.dropout, training=self.training)\n",
    "\n",
    "        x = self.lins[-1](x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Moduł PyTorch Lightning\n",
    "\n",
    "Implementujemy moduł PyTorch Lightning, który będzie zarządzał treningiem i ewaluacją modelu."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DDILinkPrediction(pl.LightningModule):\n",
    "    \"\"\"\n",
    "    Moduł PyTorch Lightning do treningu modelu przewidywania interakcji między lekami.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    in_channels : int\n",
    "        Wymiar wejściowy (liczba cech wierzchołków)\n",
    "    hidden_channels : int\n",
    "        Wymiar warstw ukrytych\n",
    "    out_channels : int\n",
    "        Wymiar embeddingów wyjściowych\n",
    "    num_layers : int\n",
    "        Liczba warstw w modelach\n",
    "    dropout : float\n",
    "        Współczynnik dropout\n",
    "    lr : float\n",
    "        Współczynnik uczenia\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, gnn, predictor, lr=0.001):\n",
    "        super().__init__()\n",
    "        self.save_hyperparameters()\n",
    "\n",
    "        # Inicjalizacja modeli\n",
    "        self.gnn = gnn\n",
    "        self.predictor = predictor\n",
    "        self.criterion = torch.nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        return self.gnn(x, edge_index)\n",
    "\n",
    "    def predict_edges(self, x, edge_index, edge_pairs):\n",
    "        \"\"\"\n",
    "        Przewiduje prawdopodobieństwo interakcji dla zadanych par wierzchołków.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : torch.Tensor\n",
    "            Macierz cech wierzchołków.\n",
    "        edge_index : torch.Tensor\n",
    "            Indeksy krawędzi w grafie.\n",
    "        edge_pairs : torch.Tensor\n",
    "            Pary wierzchołków (shape: [2, num_pairs]) dla których należy przewidzieć interakcje.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Tensor zawierający prawdopodobieństwa interakcji dla zadanych par wierzchołków.\n",
    "        \"\"\"\n",
    "        # Generowanie embeddingów dla wszystkich wierzchołków\n",
    "        embeddings = self(x, edge_index)\n",
    "\n",
    "        # Przewidywanie prawdopodobieństwa interakcji\n",
    "        predictions = self.predictor(\n",
    "            embeddings[edge_pairs[0]], embeddings[edge_pairs[1]]\n",
    "        )\n",
    "\n",
    "        return predictions\n",
    "\n",
    "    def _step(self, batch, stage):\n",
    "        \"\"\"\n",
    "        Wspólna funkcja dla kroków treningu i walidacji.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        batch : Data\n",
    "            Batch danych\n",
    "        stage : str\n",
    "            Etap ('train' lub 'val')\n",
    "        x : torch.Tensor\n",
    "            Cechy wierzchołków dla danego etapu\n",
    "        \"\"\"\n",
    "        # Generowanie embeddingów\n",
    "        embeddings = self(batch.x, batch.edge_index)\n",
    "\n",
    "        # Przewidywanie dla wszystkich krawędzi\n",
    "        pred = self.predictor(\n",
    "            embeddings[batch.edge_label_index[0]], embeddings[batch.edge_label_index[1]]\n",
    "        ).squeeze()\n",
    "\n",
    "        # Obliczanie straty\n",
    "        loss = self.criterion(pred, batch.edge_label)\n",
    "\n",
    "        # Obliczanie metryk\n",
    "        stage_auc = tm.functional.auroc(pred, batch.edge_label.long(), task='binary')\n",
    "        stage_f1 = tm.functional.f1_score(pred, batch.edge_label.long(), num_classes=2, average='micro', task='binary')\n",
    "        \n",
    "        # Logowanie metryk\n",
    "        self.log(f\"{stage}_loss\", loss, batch_size=batch.edge_label_index.size(1))\n",
    "        self.log(\n",
    "            f\"{stage}_auc\",\n",
    "            stage_auc,\n",
    "            on_step=False,\n",
    "            on_epoch=True,\n",
    "            batch_size=batch.edge_label_index.size(1),\n",
    "        )\n",
    "        self.log(\n",
    "            f\"{stage}_f1\",\n",
    "            stage_f1,\n",
    "            on_step=False,\n",
    "            on_epoch=True,\n",
    "            batch_size=batch.edge_label_index.size(1),\n",
    "        )\n",
    "\n",
    "        # Logowanie hits@k dla wszystkich etapów\n",
    "        hits_at3 = tm.functional.retrieval_recall(pred, batch.edge_label, 3)\n",
    "        hits_at5 = tm.functional.retrieval_recall(pred, batch.edge_label, 5)\n",
    "        hits_at10 = tm.functional.retrieval_recall(pred, batch.edge_label, 10)\n",
    "        hits_at20 = tm.functional.retrieval_recall(pred, batch.edge_label, 20)\n",
    "        mrr = tm.functional.retrieval_reciprocal_rank(pred, batch.edge_label)\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        return self._step(batch, \"train\")\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        return self._step(batch, \"val\")\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        return self._step(batch, \"test\")\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr)\n",
    "        return optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optymalizacja i trening modelu\n",
    "\n",
    "Trenujemy model na przygotowanych danych. Istnieje możliwość:\n",
    "1. Optymalizacji hiperparametrów;\n",
    "2. Wykorzystania najlepszych hiperparametrów do treningu modelu.\n",
    "\n",
    "\n",
    "<div class=\"alert alert-warning\">\n",
    "Uwaga, poniższe operacje mogą zająć więcej czasu - czas wykonywania jest zależny od wydajności sprzętu, oraz\n",
    "konfiguracji - wykorzystanie CPU lub GPU.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optymalizacja hiperparametrów"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model(params):\n",
    "    \"\"\"\n",
    "    Parameters\n",
    "    ----------\n",
    "    params : dict\n",
    "        Słownik zawierający hiperparametry.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    tuple\n",
    "        Krotka zawierająca model GNN oraz predyktor.\n",
    "    \"\"\"\n",
    "    hidden_channels = params['hidden_channels']\n",
    "    num_layers = params['num_layers']\n",
    "    batch_norm = params['batch_norm']\n",
    "    dropout = params['dropout']\n",
    "    selected_aggr = params['aggregation']\n",
    "    normalize = params['normalize']\n",
    "    project = params['project']\n",
    "    jk = params['jk']\n",
    "\n",
    "    # Tworzenie modelu GNN\n",
    "    gnn_layers = []\n",
    "\n",
    "    # Dodanie warstwy BatchNorm tylko jeśli batch_norm=True\n",
    "    if batch_norm:\n",
    "        gnn_layers.append((pyg.nn.BatchNorm(6), \"x -> x\"))\n",
    "\n",
    "    pyg.seed_everything(123)\n",
    "    gnn_layers.append((pyg.nn.GraphSAGE(\n",
    "        in_channels=6, \n",
    "        hidden_channels=hidden_channels, \n",
    "        out_channels=hidden_channels,\n",
    "        num_layers=num_layers,\n",
    "        dropout=dropout,  # Użycie zoptymalizowanego dropout\n",
    "        aggr=selected_aggr,\n",
    "        jk=jk,\n",
    "        normalize=normalize,\n",
    "        project=project\n",
    "    ), 'x, edge_index -> x'))\n",
    "\n",
    "    gnn = pyg.nn.Sequential(\"x, edge_index\", gnn_layers)\n",
    "\n",
    "    # Liczba warstw w predyktorze\n",
    "    n_predictor_layers = params['predictor_layers']\n",
    "    predictor_dims = params['predictor_dims']\n",
    "    predictor_dropout = params['predictor_dropout']\n",
    "    pyg.seed_everything(123)\n",
    "    predictor = LinkPredictor(hidden_channels, predictor_dims, n_predictor_layers, predictor_dropout)\n",
    "\n",
    "    return gnn, predictor\n",
    "\n",
    "def suggest_model(trial):\n",
    "    \"\"\"\n",
    "    Parameters\n",
    "    ----------\n",
    "    trial : optuna.Trial\n",
    "        Obiekt próby optymalizacji Optuna.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    tuple\n",
    "        Krotka zawierająca model GNN oraz predyktor.\n",
    "    \"\"\"\n",
    "    params = {\n",
    "        'hidden_channels': trial.suggest_categorical('hidden_channels', [32, 64, 128, 256]),\n",
    "        'num_layers': trial.suggest_int('num_layers', 3, 6),\n",
    "        'batch_norm': trial.suggest_categorical('batch_norm', [True, False]),\n",
    "        'dropout': trial.suggest_categorical('dropout', [0.1, 0.2, 0.3, 0.5]),\n",
    "        'aggregation': trial.suggest_categorical('aggregation', [\n",
    "            'add', 'mean', 'max',\n",
    "            ('add', 'mean'),\n",
    "            ('add', 'max'),\n",
    "            ('mean', 'max'),\n",
    "            ('add', 'mean', 'max')\n",
    "        ]),\n",
    "        'normalize': trial.suggest_categorical('normalize', [True, False]),\n",
    "        'project': trial.suggest_categorical('project', [True, False]),\n",
    "        'jk': trial.suggest_categorical('jk', [None, 'last', 'cat', 'lstm']),\n",
    "        'predictor_layers': trial.suggest_int('predictor_layers', 2, 4),\n",
    "        'predictor_dims': trial.suggest_categorical('predictor_dims', [32, 64, 128, 256]),\n",
    "        'predictor_dropout': trial.suggest_categorical('predictor_dropout', [0.1, 0.2, 0.3, 0.5])\n",
    "    }\n",
    "\n",
    "    return create_model(params)\n",
    "\n",
    "def objective(trial):\n",
    "    \"\"\"\n",
    "    Parameters\n",
    "    ----------\n",
    "    trial : optuna.Trial\n",
    "        Obiekt próby optymalizacji Optuna.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        Wartość metryki walidacyjnej (AUC).\n",
    "    \"\"\"\n",
    "    gnn, predictor = suggest_model(trial)\n",
    "    model = DDILinkPrediction(gnn, predictor)\n",
    "    \n",
    "    trainer = pl.Trainer(\n",
    "        max_epochs=5,\n",
    "        accelerator=\"gpu\" if torch.cuda.is_available() else \"cpu\",\n",
    "        devices=1,\n",
    "        enable_checkpointing=False,\n",
    "        logger=False,\n",
    "        enable_progress_bar=False\n",
    "    )\n",
    "    \n",
    "    trainer.fit(model, train_loader, val_loader)\n",
    "    val_results = trainer.validate(model, val_loader)\n",
    "    \n",
    "    return val_results[0]['val_f1']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "tune = False\n",
    "if tune:\n",
    "    # Utworzenie studium Optuna\n",
    "    study = optuna.create_study(direction='maximize')\n",
    "    study.optimize(objective, n_trials=3, timeout=3600/4) # 15 min. max\n",
    "    # Użycie najlepszych hiperparametrów\n",
    "    with open(\"linkpred.json\", \"w\") as f:\n",
    "        json.dump(study.best_params, f)\n",
    "    best_params = study.best_params\n",
    "else:\n",
    "    with open(\"linkpred.json\", \"r\") as f:\n",
    "        best_params = json.load(f)\n",
    "gnn, predictor = create_model(best_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'hidden_channels': 128,\n",
       " 'num_layers': 3,\n",
       " 'batch_norm': True,\n",
       " 'dropout': 0.2,\n",
       " 'aggregation': ['add', 'max', 'mean'],\n",
       " 'normalize': True,\n",
       " 'project': True,\n",
       " 'jk': 'cat',\n",
       " 'predictor_layers': 3,\n",
       " 'predictor_dims': 256,\n",
       " 'predictor_dropout': 0.5}"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg.seed_everything(123)\n",
    "# Inicjalizacja modelu\n",
    "gnn, predictor = create_model(best_params)\n",
    "model = DDILinkPrediction(gnn, predictor)\n",
    "\n",
    "# Inicjalizacja trenera\n",
    "trainer = pl.Trainer(\n",
    "    max_epochs=25,\n",
    "    accelerator=\"gpu\" if torch.cuda.is_available() else \"cpu\",\n",
    "    devices=1,\n",
    "    deterministic=True,\n",
    "    log_every_n_steps=1,\n",
    "    callbacks=[\n",
    "        pl.callbacks.ModelCheckpoint(monitor=\"val_f1\", mode=\"max\"),\n",
    "        pl.callbacks.EarlyStopping(monitor=\"val_f1\", mode=\"max\", patience=5),\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Wczytanie najlepszego modelu z punktu kontrolnego\n",
    "best_model_path = trainer.checkpoint_callback.best_model_path\n",
    "if best_model_path:\n",
    "    print(f\"Wczytywanie najlepszego modelu z: {best_model_path}\")\n",
    "    model = DDILinkPrediction.load_from_checkpoint(\n",
    "        best_model_path, gnn=gnn, predictor=predictor\n",
    "    )\n",
    "else:\n",
    "    print(\"Nie znaleziono punktu kontrolnego - użycie aktualnego modelu\")\n",
    "\n",
    "\n",
    "# Walidacja modelu przed treningiem\n",
    "print(\"Walidacja modelu przed treningiem:\")\n",
    "val_scores_before_train = trainer.validate(model, val_loader)\n",
    "\n",
    "# Trenowanie modelu\n",
    "trainer.fit(model, train_loader, val_loader)\n",
    "\n",
    "# Walidacja modelu po treningu\n",
    "print(\"Walidacja modelu po treningu:\")\n",
    "val_scores_after_train = trainer.validate(model, val_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "/home/filip/projects/grafowe_sieci_neuronowe/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:433: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9c2e46c66db24465902f61a0e4b84a3a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "        test_auc            0.9153414368629456\n",
      "         test_f1            0.8541463017463684\n",
      "        test_loss           0.36219704151153564\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 0.36219704151153564,\n",
       "  'test_auc': 0.9153414368629456,\n",
       "  'test_f1': 0.8541463017463684}]"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Sprawdzenie na zbiorze testowym\n",
    "test_scores_after = trainer.test(model, test_loader)\n",
    "test_scores_after"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ewaluacja top K"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TorchMetrics udostępnia dedykowane funkcje do obliczania metryk w zakresie TOPK (Hits@K), takie jak RetrievalRecall@K, które są tutaj wykorzystywane do ewaluacji skuteczności modelu dla różnych wartości K."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Wyniki dla zbioru walidacyjnego:\n",
      "--------------------------------------------------\n",
      "\n",
      "Metryki Hits@K:\n",
      "Hits@1: 0.0905\n",
      "Hits@3: 0.2331\n",
      "Hits@5: 0.3477\n",
      "Hits@10: 0.5783\n",
      "Hits@20: 0.8013\n",
      "Hits@50: 0.8803\n",
      "\n",
      "Wyniki dla zbioru testowego:\n",
      "--------------------------------------------------\n",
      "\n",
      "Metryki Hits@K:\n",
      "Hits@1: 0.0862\n",
      "Hits@3: 0.2169\n",
      "Hits@5: 0.3346\n",
      "Hits@10: 0.5628\n",
      "Hits@20: 0.7856\n",
      "Hits@50: 0.8639\n"
     ]
    }
   ],
   "source": [
    "def evaluate_model(model, data, eval_data, split_name=\"\"):\n",
    "    \"\"\"\n",
    "    Parameters\n",
    "    ----------\n",
    "    model : torch.nn.Module\n",
    "        Model do ewaluacji.\n",
    "    data : Data\n",
    "        Dane wejściowe zawierające cechy wierzchołków.\n",
    "    eval_data : Data\n",
    "        Dane do ewaluacji (walidacyjne lub testowe).\n",
    "    split_name : str\n",
    "        Nazwa zbioru danych (np. \"Walidacyjny\" lub \"Testowy\").\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        yhat = torch.sigmoid(\n",
    "            model.predict_edges(\n",
    "                data.x, eval_data.edge_label_index, eval_data.edge_label_index\n",
    "            )\n",
    "        )\n",
    "\n",
    "    print(f\"\\nWyniki dla zbioru {split_name}:\")\n",
    "    print(\"-\" * 50)\n",
    "\n",
    "    print(\"\\nMetryki Hits@K:\")\n",
    "    for k in [1, 3, 5, 10, 20, 50]:\n",
    "        hit_rate = tm.retrieval.RetrievalRecall(top_k=k)\n",
    "        score = hit_rate(\n",
    "            yhat.squeeze(),\n",
    "            eval_data.edge_label,\n",
    "            eval_data.edge_label_index[0, :],\n",
    "        )\n",
    "        print(f\"Hits@{k}: {score:.4f}\")\n",
    "    \n",
    "\n",
    "\n",
    "# Ewaluacja na zbiorze walidacyjnym\n",
    "evaluate_model(model, val_data, val_data, \"walidacyjnego\")\n",
    "\n",
    "# Ewaluacja na zbiorze testowym\n",
    "evaluate_model(model, test_data, test_data, \"testowego\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gnnbook (3.13.5)",
   "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.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
