{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wprowadzenie"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Niniejszy notatnik zawiera przykład klasyfikacji całych grafów na przykładzie popularnego zbioru danych MUTAG, używanego niejednokrotnie jako benchmark dla nowych modeli GNN.\n",
    "\n",
    "Składa się on ze 188 grafów, które reprezentują różne związki chemiczne. Każdy graf w zbiorze odpowiada pojedynczej cząsteczce, gdzie wierzchołki reprezentują atomy, a krawędzie odpowiadają wiązaniom chemicznym między nimi. Celem analizy jest klasyfikacja tych cząsteczek do jednej z dwóch klas: mutagenicznych (około 66.5% zbioru) lub niemutagenicznych (około 33.5%), co odnosi się do ich potencjalnego wpływu na mutacje genetyczne w organizmach, badanych na przykładzie bakterii Salmonella typhimurium."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import bibliotek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch_geometric as pyg\n",
    "import torch as th\n",
    "import torch.nn.functional as F\n",
    "import pytorch_lightning as pl\n",
    "import networkx as nx\n",
    "import pandas as pd\n",
    "import optuna as opt\n",
    "import torchmetrics as tm\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as mpatches\n",
    "\n",
    "from typing import Dict, Any, Literal\n",
    "from sklearn.metrics import classification_report\n",
    "from sklearn.model_selection import KFold\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "import utils.vis_utils as vis_utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przykład budowania próbek uczących"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 29132 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "g1 = pyg.data.Data()\n",
    "g1.x = th.tensor([[i, i] for i in range(1, 6)], dtype=th.float)\n",
    "g1.edge_index = th.tensor([\n",
    "    [0, 1, 1, 2, 2, 3],\n",
    "    [1, 2, 3, 3, 4, 4]\n",
    "])\n",
    "\n",
    "g2 = pyg.data.Data()\n",
    "g2.x = th.tensor([[i, i] for i in range(6, 10)], dtype=th.float)\n",
    "g2.edge_index = th.tensor([\n",
    "    [0, 2, 3],\n",
    "    [1, 1, 1]   \n",
    "])\n",
    "\n",
    "g1_nx = pyg.utils.to_networkx(g1)\n",
    "pyg.seed_everything(123)\n",
    "pos = nx.spring_layout(g1_nx)\n",
    "nx.draw(g1_nx, with_labels=True, pos=pos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 24660 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "g2_nx = pyg.utils.to_networkx(g2)\n",
    "pyg.seed_everything(123)\n",
    "pos = nx.spring_layout(g2_nx)\n",
    "nx.draw(g2_nx, with_labels=True, pos=pos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "small_data_loader = pyg.loader.DataLoader([g1, g2], batch_size=10)\n",
    "batch = next(iter(small_data_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DataBatch(x=[9, 2], edge_index=[2, 9], batch=[9], ptr=[3])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 1.],\n",
       "        [2., 2.],\n",
       "        [3., 3.],\n",
       "        [4., 4.],\n",
       "        [5., 5.],\n",
       "        [6., 6.],\n",
       "        [7., 7.],\n",
       "        [8., 8.],\n",
       "        [9., 9.]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch.x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 1, 1, 2, 2, 3, 5, 7, 8],\n",
       "        [1, 2, 3, 3, 4, 4, 6, 6, 6]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch.edge_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 0, 0, 0, 0, 1, 1, 1, 1])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch.batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przygotowanie danych"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pobranie zbioru"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://www.chrsmrrs.com/graphkerneldatasets/MUTAG.zip\n",
      "Processing...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "dataset = pyg.datasets.TUDataset(root='data/TUDataset', name='MUTAG')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ilość grafów: 188\n",
      "Ilość cech w grafie: 7\n",
      "Ilość klas: 2\n",
      "Proporcje klas: [0.66489362 0.33510638]\n"
     ]
    }
   ],
   "source": [
    "# Ilość grafów w zbiorze danych\n",
    "num_graphs = len(dataset)\n",
    "\n",
    "# Ilość cech w grafie\n",
    "num_features = dataset.num_features\n",
    "\n",
    "# Ilość klas\n",
    "num_classes = dataset.num_classes\n",
    "\n",
    "class_props = pd.Series(dataset.y).value_counts(normalize=True)\n",
    "\n",
    "print(f'Ilość grafów: {num_graphs}')\n",
    "print(f'Ilość cech w grafie: {num_features}')\n",
    "print(f'Ilość klas: {num_classes}')\n",
    "print(f'Proporcje klas: {class_props.values}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Średni stopień wierzchołka: 4.470588235294118\n",
      "Czy graf jest skierowany: True\n",
      "Ilość wierzchołków: 17\n",
      "Ilość krawędzi: 38\n"
     ]
    }
   ],
   "source": [
    "# Pobierz pierwszy graf ze zbioru\n",
    "data = dataset[0]\n",
    "\n",
    "# Konwersja do grafu NetworkX\n",
    "G = pyg.utils.to_networkx(data)\n",
    "\n",
    "# Średni stopień wierzchołka\n",
    "average_degree = sum(dict(G.degree()).values()) / G.number_of_nodes()\n",
    "\n",
    "# Czy graf jest skierowany\n",
    "is_directed = nx.is_directed(G)\n",
    "\n",
    "# Ilość wierzchołków\n",
    "num_nodes = G.number_of_nodes()\n",
    "\n",
    "# Ilość krawędzi\n",
    "num_edges = G.number_of_edges()\n",
    "\n",
    "print(f'Średni stopień wierzchołka: {average_degree}')\n",
    "print(f'Czy graf jest skierowany: {is_directed}')\n",
    "print(f'Ilość wierzchołków: {num_nodes}')\n",
    "print(f'Ilość krawędzi: {num_edges}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Podział danych"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg.seed_everything(1234)\n",
    "dataset = dataset.shuffle()\n",
    "dataset.y = dataset.y.to(th.long)\n",
    "n_train = int(0.7 * num_graphs)\n",
    "n_test = num_graphs - n_train\n",
    "\n",
    "train_dataset = dataset[:n_train]\n",
    "test_dataset = dataset[n_train:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(131, 57)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n_train, n_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_train_loader = pyg.loader.DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
    "test_loader = pyg.loader.DataLoader(test_dataset, batch_size=40, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_b = next(iter(all_train_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>1</th>\n",
       "      <th>0</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>proportion</th>\n",
       "      <td>0.781250</td>\n",
       "      <td>0.218750</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>proportion</th>\n",
       "      <td>0.625000</td>\n",
       "      <td>0.375000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>proportion</th>\n",
       "      <td>0.687500</td>\n",
       "      <td>0.312500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>proportion</th>\n",
       "      <td>0.687500</td>\n",
       "      <td>0.312500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>proportion</th>\n",
       "      <td>0.666667</td>\n",
       "      <td>0.333333</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                   1         0\n",
       "proportion  0.781250  0.218750\n",
       "proportion  0.625000  0.375000\n",
       "proportion  0.687500  0.312500\n",
       "proportion  0.687500  0.312500\n",
       "proportion  0.666667  0.333333"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pos_neg_probs = []\n",
    "for b in all_train_loader:\n",
    "    pos_neg_probs.append(pd.Series(b.y).value_counts(normalize=True))\n",
    "    \n",
    "pos_neg_probs_df = pd.DataFrame(pos_neg_probs)\n",
    "pos_neg_probs_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przygotowanie modelu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Słownik z nazwami modeli i odpowiadającymi im klasami\n",
    "model_names2class: Dict[str, Any] = {\n",
    "    'GCN': pyg.nn.GCN,\n",
    "    'SAGE': pyg.nn.GraphSAGE,\n",
    "    'GAT': pyg.nn.GAT,\n",
    "    \"GIN\": pyg.nn.GIN\n",
    "}\n",
    "\n",
    "class GraphClassifier(th.nn.Module):\n",
    "    \"\"\"\n",
    "    Klasyfikator grafów\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    model_name : str\n",
    "        Nazwa modelu do użycia ('GCN', 'SAGE', 'GAT').\n",
    "    gnn_num_layers : int\n",
    "        Liczba warstw w modelu.\n",
    "    gnn_hidden_dim : int\n",
    "        Wymiar ukrytej warstwy.\n",
    "    gnn_input_dim : int\n",
    "        Wymiar wejściowy.\n",
    "    num_classes : int\n",
    "        Liczba klas wyjściowych.\n",
    "    dropout : float, optional\n",
    "        Wartość dropout (domyślnie 0.2).\n",
    "    **kwargs : dict\n",
    "        Dodatkowe argumenty przekazywane do konstruktora modelu.\n",
    "\n",
    "    Attributes\n",
    "    ----------\n",
    "    gnn_model : torch_geometric.nn.MessagePassing\n",
    "        Zbudowany model grafowej sieci neuronowej.\n",
    "    classifier : torch.nn.Sequential\n",
    "        Warstwa klasyfikatora.\n",
    "    \"\"\"\n",
    "    def __init__(self, \n",
    "                 model_name: str, \n",
    "                 gnn_num_layers: int, \n",
    "                 gnn_hidden_dim: int,\n",
    "                 gnn_input_dim: int,\n",
    "                 num_classes: int,\n",
    "                 lin_sizes: list[int] = [],\n",
    "                 dropout: float = 0.2,\n",
    "                 pool_type: Literal['add', 'mean'] = 'add',\n",
    "                 **kwargs: Any) -> None:\n",
    "        super().__init__()\n",
    "        self.model_name = model_name\n",
    "        self.num_layers = gnn_num_layers\n",
    "        self.hidden_dim = gnn_hidden_dim\n",
    "        self.input_dim = gnn_input_dim\n",
    "        self.num_classes = num_classes\n",
    "        self.dropout = dropout\n",
    "        self.pool_type = pool_type\n",
    "        self.gnn_model = self.build_gnn_model(**kwargs)\n",
    "        self.pool_layer = pyg.nn.global_add_pool if pool_type == 'add' else pyg.nn.global_mean_pool\n",
    "        all_lin_sizes = [self.hidden_dim] + lin_sizes + [num_classes]\n",
    "        lin_layers = []\n",
    "        for i, (dim_in, dim_out) in enumerate(zip(all_lin_sizes, all_lin_sizes[1:])):\n",
    "            lin_layers.append(th.nn.Dropout(dropout))\n",
    "            lin_layers.append(th.nn.Linear(dim_in, dim_out))\n",
    "            if i < len(all_lin_sizes) - 2:\n",
    "                lin_layers.append(th.nn.ReLU())\n",
    "        self.classifier = th.nn.Sequential(*lin_layers)\n",
    "\n",
    "    def build_gnn_model(self, **kwargs) -> pyg.nn.MessagePassing:\n",
    "        \"\"\"\n",
    "        Buduje model na podstawie zadanej architektury.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        **kwargs : dict\n",
    "            Dodatkowe argumenty przekazywane do konstruktora modelu.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        model : pyg.nn.MessagePassing\n",
    "            Zbudowany model grafowej sieci neuronowej.\n",
    "        \"\"\"\n",
    "        model_arch = model_names2class[self.model_name]\n",
    "        model = model_arch(\n",
    "            in_channels=self.input_dim,\n",
    "            hidden_channels=self.hidden_dim,\n",
    "            num_layers=self.num_layers,\n",
    "            **kwargs\n",
    "        )\n",
    "        return model\n",
    "        \n",
    "    def forward(self, x: th.Tensor, edge_index: th.Tensor, batch_idx: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"\n",
    "        Przejście przez warstwy modelu.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : torch.Tensor\n",
    "            Macierz cech wierzchołków.\n",
    "        edge_index : torch.Tensor\n",
    "            Macierz indeksów krawędzi.\n",
    "        batch_idx : torch.Tensor\n",
    "            Indeksy wsadowe dla grafów w partii.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        yhat : torch.Tensor\n",
    "            Wektor predykcji po przejściu przez model.\n",
    "        \"\"\"\n",
    "        # Przejście przez model GNN\n",
    "        embed = self.gnn_model(x, edge_index)\n",
    "        \n",
    "        # Globalne zagregowanie osadzeń wierzchołków dla każdego grafu w partii\n",
    "        # Funkcja agreguje osadzenia wierzchołków do jednego wektora reprezentującego cały graf\n",
    "        mean_pool = self.pool_layer(embed, batch_idx)\n",
    "        \n",
    "        # Przejście przez warstwę klasyfikatora\n",
    "        yhat = self.classifier(mean_pool)\n",
    "        return yhat\n",
    "    \n",
    "\n",
    "class GraphClassifierPL(pl.LightningModule):\n",
    "    def __init__(self, model: GraphClassifier, lr: float = 0.001):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "        self.lr = lr\n",
    "        self.criterion = th.nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, x, edge_index, batch_idx):\n",
    "        return self.model(x, edge_index, batch_idx)\n",
    "\n",
    "    def common_step(self, batch, prefix):\n",
    "        x, edge_index, batch_idx, y = batch.x, batch.edge_index, batch.batch, batch.y\n",
    "        y_hat = self(x, edge_index, batch_idx)\n",
    "        loss = self.criterion(y_hat, y.to(th.long))\n",
    "        batch_size = y.size(0)\n",
    "        \n",
    "        f1 = tm.functional.f1_score(y_hat, y, task='multiclass', average='weighted', num_classes=self.model.num_classes)\n",
    "        acc = tm.functional.accuracy(y_hat, y, task='multiclass', average='weighted', num_classes=self.model.num_classes)\n",
    "        recall = tm.functional.recall(y_hat, y, task='multiclass', average='weighted', num_classes=self.model.num_classes)\n",
    "        \n",
    "        self.log(f\"{prefix}_loss\", loss, on_step=False, on_epoch=True, batch_size=batch_size, prog_bar=True)\n",
    "        self.log(f\"{prefix}_f1\", f1, on_step=False, on_epoch=True, batch_size=batch_size, prog_bar=True)\n",
    "        self.log(f\"{prefix}_acc\", acc, on_step=False, on_epoch=True, batch_size=batch_size, prog_bar=True)\n",
    "        self.log(f\"{prefix}_recall\", recall, on_step=False, on_epoch=True, batch_size=batch_size, prog_bar=True)\n",
    "        \n",
    "        return loss\n",
    "\n",
    "    def training_step(self, batch):\n",
    "        return self.common_step(batch, \"train\")\n",
    "\n",
    "    def validation_step(self, batch):\n",
    "        return self.common_step(batch, \"val\")\n",
    "\n",
    "    def test_step(self, batch):\n",
    "        return self.common_step(batch, \"test\")\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return th.optim.Adam(self.parameters(), lr=self.lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/filip/projects/grafowe_sieci_neuronowe/.venv/lib/python3.12/site-packages/pytorch_lightning/core/module.py:449: You are trying to `self.log()` but the `self.trainer` reference is not registered on the model yet. This is most likely because the model hasn't been passed to the `Trainer`\n"
     ]
    }
   ],
   "source": [
    "# Przykład użycia:\n",
    "pyg.seed_everything(123)\n",
    "model_params = {'model_name': 'GIN', 'gnn_num_layers': 4, 'gnn_hidden_dim': 32, 'gnn_input_dim': 7, 'lin_sizes': [16, 8], 'num_classes': 2, 'dropout': 0.2, 'aggr': ['add'], 'jk': 'cat'}\n",
    "#model = GraphClassifier(model_name='GAT', heads=4, gnn_num_layers=4, gnn_hidden_dim=64, gnn_input_dim=num_features, lin_sizes=[16, 32], num_classes=1, aggr=['sum'], jk='cat')\n",
    "model = GraphClassifier(**model_params)\n",
    "model_lit = GraphClassifierPL(model=model, lr=0.001)\n",
    "with th.no_grad():\n",
    "    out = model(train_b.x, train_b.edge_index, train_b.batch)\n",
    "    \n",
    "    out_lit = model_lit(train_b.x, train_b.edge_index, train_b.batch)\n",
    "    \n",
    "    model_lit.training_step(train_b)\n",
    "    model_lit.validation_step(train_b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GraphClassifier(\n",
       "  (gnn_model): GIN(7, 32, num_layers=4)\n",
       "  (classifier): Sequential(\n",
       "    (0): Dropout(p=0.2, inplace=False)\n",
       "    (1): Linear(in_features=32, out_features=16, bias=True)\n",
       "    (2): ReLU()\n",
       "    (3): Dropout(p=0.2, inplace=False)\n",
       "    (4): Linear(in_features=16, out_features=8, bias=True)\n",
       "    (5): ReLU()\n",
       "    (6): Dropout(p=0.2, inplace=False)\n",
       "    (7): Linear(in_features=8, out_features=2, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optymalizacja hiperparametrów"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def optuna_params_to_model_params(optuna_params: dict) -> dict:\n",
    "    \"\"\"\n",
    "    Konwersja słownika z optymalnymi parametrami z Optuna na słownik parametrów modelu.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    optuna_params : dict\n",
    "        Słownik z optymalnymi parametrami z Optuna.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    model_params : dict\n",
    "        Słownik z parametrami modelu.\n",
    "    \"\"\"\n",
    "    model_params = {\n",
    "        \"model_name\": optuna_params[\"conv_layer\"],\n",
    "        \"gnn_num_layers\": optuna_params[\"gnn_num_layers\"],\n",
    "        \"gnn_hidden_dim\": optuna_params[\"gnn_hidden_dim\"],\n",
    "        \"gnn_input_dim\": num_features,\n",
    "        \"lin_sizes\": [optuna_params[\"lin_size\"]] * optuna_params[\"num_lins\"],\n",
    "        \"num_classes\": num_classes,\n",
    "        \"dropout\": 0.2,\n",
    "        \"pool_type\": optuna_params.get(\"pool_type\"),\n",
    "        \"aggr\": optuna_params.get(\"aggr_multi\", optuna_params.get(\"aggr_single\")),\n",
    "        \"jk\": optuna_params.get(\"jk\"),\n",
    "    }\n",
    "    if optuna_params[\"conv_layer\"] == \"SAGE\":\n",
    "        model_params[\"normalize\"] = optuna_params.get(\"normalize\")\n",
    "        model_params[\"project\"] = optuna_params.get(\"project\")\n",
    "    elif optuna_params[\"conv_layer\"] == \"GAT\":\n",
    "        model_params[\"heads\"] = optuna_params.get(\"num_heads\")\n",
    "        model_params[\"dropout\"] = optuna_params.get(\"dropout\")\n",
    "    elif optuna_params[\"conv_layer\"] == \"GCN\":\n",
    "        model_params[\"normalize\"] = optuna_params.get(\"normalize\")\n",
    "    return model_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_proposed_params(trial) -> dict:\n",
    "    conv_layer_name = trial.suggest_categorical(\n",
    "        \"conv_layer\", list(model_names2class.keys())\n",
    "    )\n",
    "    \n",
    "    gnn_num_layers = trial.suggest_int(\"gnn_num_layers\", 1, 3)\n",
    "    gnn_hidden_dim = trial.suggest_categorical(\"gnn_hidden_dim\", [8, 16, 32, 64, 128])\n",
    "    jk = trial.suggest_categorical(\"jk\", [\"cat\", \"max\", \"lstm\", \"last\", None])\n",
    "    pool_type = trial.suggest_categorical(\"pool_type\", [\"add\", \"mean\"])\n",
    "    num_lins = trial.suggest_int(\"num_lins\", 0, 2)\n",
    "    lin_size = trial.suggest_categorical(\"lin_size\", [8, 16, 32])\n",
    "    model_kwargs = {}\n",
    "    if conv_layer_name == \"SAGE\":\n",
    "        aggr = trial.suggest_categorical(\n",
    "            \"aggr_multi\",\n",
    "            [\n",
    "                \"sum\",\n",
    "                \"mean\",\n",
    "                \"max\",\n",
    "                [\"mean\", \"sum\"],\n",
    "                [\"sum\", \"max\"],\n",
    "                [\"sum\", \"max\", \"mean\"],\n",
    "            ],\n",
    "        )\n",
    "        normalize = trial.suggest_categorical(\"normalize\", [True, False])\n",
    "        project = trial.suggest_categorical(\"project\", [True, False])\n",
    "        model_kwargs[\"normalize\"] = normalize\n",
    "        model_kwargs[\"project\"] = project\n",
    "    elif conv_layer_name == \"GAT\":\n",
    "        aggr = trial.suggest_categorical(\"aggr_single\", [\"sum\", \"mean\", \"max\"])\n",
    "        num_heads = trial.suggest_categorical(\"num_heads\", [1, 2])\n",
    "        dropout = trial.suggest_categorical(\"dropout\", [0, 0.2, 0.5])\n",
    "        model_kwargs[\"heads\"] = num_heads\n",
    "        model_kwargs[\"dropout\"] = dropout\n",
    "    elif conv_layer_name == \"GCN\":\n",
    "        aggr = trial.suggest_categorical(\"aggr_single\", [\"sum\", \"mean\", \"max\"])\n",
    "        normalize = trial.suggest_categorical(\"normalize\", [True, False])\n",
    "        model_kwargs[\"normalize\"] = normalize\n",
    "    elif conv_layer_name == \"GIN\":\n",
    "        aggr = trial.suggest_categorical(\"aggr_single\", [\"sum\", \"mean\", \"max\"])\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown convolution layer: {conv_layer_name}\")\n",
    "\n",
    "    params_dict = {\n",
    "        \"model_name\": conv_layer_name,\n",
    "        \"gnn_num_layers\": gnn_num_layers,\n",
    "        \"gnn_hidden_dim\": gnn_hidden_dim,\n",
    "        \"gnn_input_dim\": num_features,\n",
    "        \"pool_type\": pool_type,\n",
    "        \"lin_sizes\": [lin_size] * num_lins,\n",
    "        \"num_classes\": num_classes,\n",
    "        \"aggr\": aggr,\n",
    "        \"jk\": jk,\n",
    "    } | model_kwargs\n",
    "    return params_dict\n",
    "    \n",
    "    \n",
    "    \n",
    "def objective(trial):\n",
    "    \"\"\"\n",
    "    Funkcja celu dla procesu optymalizacji hiperparametrów przy użyciu biblioteki Optuna.\n",
    "\n",
    "    Funkcja przeprowadza proces strojenia hiperparametrów dla modelu sieci neuronowej GNN,\n",
    "    testując różne konfiguracje warstw konwolucyjnych, metod agregacji oraz architektury sieci.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    trial : optuna.trial.Trial\n",
    "        Obiekt próby Optuna, wykorzystywany do sugerowania wartości hiperparametrów.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        Wartość funkcji celu - wynik F1 na zbiorze walidacyjnym dla danej konfiguracji hiperparametrów.\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "    Funkcja optymalizuje następujące hiperparametry:\n",
    "    - Typ warstwy konwolucyjnej (GCNConv, SAGEConv, GATv2Conv)\n",
    "    - Metoda agregacji (sum, mean, max lub ich kombinacje dla SAGEConv)\n",
    "    - Liczba warstw (2-4)\n",
    "    - Rozmiary warstw ukrytych\n",
    "    - Liczba głów uwagi (tylko dla GATv2Conv)\n",
    "    \"\"\"\n",
    "    \n",
    "    proposed_params = get_proposed_params(trial)\n",
    "    kfold = KFold(n_splits=5, shuffle=True, random_state=123)\n",
    "    val_f1_scores = []\n",
    "    graph_indcices = list(range(len(train_dataset)))\n",
    "    for train_idx, val_idx in kfold.split(graph_indcices):\n",
    "        cv_train_loader = pyg.data.DataLoader([train_dataset[i] for i in train_idx], batch_size=64, shuffle=True)\n",
    "        cv_val_loader = pyg.data.DataLoader([train_dataset[i] for i in val_idx], batch_size=128, shuffle=False)\n",
    "        \n",
    "        pyg.seed_everything(123)\n",
    "        model = GraphClassifier(**proposed_params)\n",
    "        model_lit = GraphClassifierPL(model=model, lr=0.005)\n",
    "        \n",
    "        early_stopping = pl.callbacks.EarlyStopping(\n",
    "            monitor=\"val_f1\",\n",
    "            patience=5,\n",
    "            mode=\"max\",\n",
    "            verbose=False\n",
    "        )\n",
    "        \n",
    "        trainer = pl.Trainer(\n",
    "            max_epochs=15,\n",
    "            enable_progress_bar=False,\n",
    "            callbacks=[early_stopping],\n",
    "            enable_checkpointing=False,\n",
    "            logger=False,\n",
    "            deterministic=True,\n",
    "            accelerator='auto',\n",
    "        )\n",
    "        \n",
    "        trainer.fit(model_lit, cv_train_loader, cv_val_loader)\n",
    "        val_f1 = trainer.callback_metrics[\"val_f1\"].item()\n",
    "        val_f1_scores.append(val_f1)\n",
    "\n",
    "    mean_val_f1 = np.mean(val_f1_scores)\n",
    "    return mean_val_f1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "tune = False\n",
    "if tune:\n",
    "    # Create an Optuna study\n",
    "    study = opt.create_study(direction='maximize')\n",
    "\n",
    "    # Optimize the study using the objective function\n",
    "    study.optimize(objective, n_trials=100, timeout=60 * 10)\n",
    "\n",
    "    # Print the best trial\n",
    "    print(f'Best trial: {study.best_trial.value}')\n",
    "    print(f'Best parameters: {study.best_trial.params}')\n",
    "\n",
    "    best_model_params = optuna_params_to_model_params(study.best_trial.params)\n",
    "    with open('best_params_graph_clf.json', 'w') as f:\n",
    "        json.dump(best_model_params, f)\n",
    "else:\n",
    "    best_model_params = json.load(open('best_params_graph_clf.json', 'r'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'model_name': 'SAGE',\n",
       " 'gnn_num_layers': 4,\n",
       " 'gnn_hidden_dim': 128,\n",
       " 'gnn_input_dim': 7,\n",
       " 'lin_sizes': [],\n",
       " 'num_classes': 2,\n",
       " 'dropout': 0.2,\n",
       " 'pool_type': 'add',\n",
       " 'aggr': 'sum',\n",
       " 'jk': 'cat'}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_model_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Szkolenie modelu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg.seed_everything(123)\n",
    "model = GraphClassifier(**best_model_params)\n",
    "model_lit = GraphClassifierPL(model=model, lr=0.001)\n",
    "\n",
    "checkpoint_callback = pl.callbacks.ModelCheckpoint(\n",
    "    monitor=\"val_f1\",\n",
    "    mode=\"max\",\n",
    "    save_top_k=1,\n",
    "    save_weights_only=True,\n",
    "    dirpath=\"checkpoints\",\n",
    "    filename='{epoch}-{val_f1:.2f}'\n",
    ")\n",
    "\n",
    "trainer = pl.Trainer(\n",
    "    max_epochs=25,\n",
    "    accelerator='auto',\n",
    "    deterministic=True,\n",
    "    callbacks=[\n",
    "        pl.callbacks.EarlyStopping(monitor='val_f1', patience=10, mode='max'),\n",
    "        checkpoint_callback\n",
    "    ]\n",
    ")\n",
    "\n",
    "trainer.fit(model_lit, all_train_loader, all_train_loader)\n",
    "\n",
    "# # Load the best model\n",
    "best_model_path = checkpoint_callback.best_model_path\n",
    "best_model = GraphClassifier(**best_model_params)\n",
    "best_model_lit = GraphClassifierPL.load_from_checkpoint(best_model_path, model=best_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "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": "e244e73cabee499fab3125fa98051b5f",
       "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_acc            0.9122806787490845\n",
      "         test_f1            0.9118067622184753\n",
      "        test_loss           0.3185197710990906\n",
      "       test_recall          0.9122806787490845\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 0.3185197710990906,\n",
       "  'test_f1': 0.9118067622184753,\n",
       "  'test_acc': 0.9122806787490845,\n",
       "  'test_recall': 0.9122806787490845}]"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.test(best_model_lit, test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_model_lit.eval()\n",
    "test_b = next(iter(test_loader))\n",
    "with th.no_grad():\n",
    "    yhat_test = best_model_lit(test_b.x, test_b.edge_index, test_b.batch).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "WYNIKI NA ZBIORZE TESTOWYM:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.93      0.82      0.88        17\n",
      "           1       0.88      0.96      0.92        23\n",
      "\n",
      "    accuracy                           0.90        40\n",
      "   macro avg       0.91      0.89      0.90        40\n",
      "weighted avg       0.90      0.90      0.90        40\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\\n\\nWYNIKI NA ZBIORZE TESTOWYM:\")\n",
    "print(classification_report(test_b.y, yhat_test.argmax(-1)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wizualizacja wektorów osadzeń"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 53512 characters]",
      "text/plain": [
       "<Figure size 472.441x472.441 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "all_data_loader = pyg.loader.DataLoader(dataset, batch_size=256, shuffle=False)\n",
    "# Pobranie całego zbioru danych\n",
    "all_batch = next(iter(all_data_loader))\n",
    "\n",
    "# Uzyskanie osadzeń wierzchołków za pomocą modelu GNN\n",
    "node_embeddings = best_model.gnn_model(all_batch.x, all_batch.edge_index)\n",
    "\n",
    "# Agregacja osadzeń wierzchołków do jednego wektora reprezentującego cały graf\n",
    "# Funkcja pool_layer agreguje osadzenia wierzchołków dla każdego grafu w partii\n",
    "graph_embedding = best_model.pool_layer(node_embeddings, all_batch.batch)\n",
    "\n",
    "# Redukcja wymiarów osadzeń grafów do 2D za pomocą PCA\n",
    "pca = PCA(n_components=2)\n",
    "embeddings_2d = pca.fit_transform(graph_embedding.detach().numpy())\n",
    "\n",
    "# Ustal odcienie szarości wyraźnie różniące się dla każdej klasy\n",
    "# 0: czarny ('#000000'), 1: jasnoszary ('#bbbbbb')\n",
    "class_to_color = {0: \"#111111\", 1: \"#bbbbbb\"}\n",
    "colors = [class_to_color[int(label)] for label in all_batch.y]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=vis_utils.image_size_in_cm(12, 12))\n",
    "scatter = ax.scatter(\n",
    "    embeddings_2d[:, 0],\n",
    "    embeddings_2d[:, 1],\n",
    "    c=colors,\n",
    "    alpha=1.0,\n",
    "    edgecolors='k',\n",
    "    linewidths=0.5\n",
    ")\n",
    "plt.xlabel('Główna składowa 1', fontsize=12)\n",
    "plt.ylabel('Główna składowa 2', fontsize=12)\n",
    "# Utwórz legendę ręcznie ze względu na opcjonalne kolory\n",
    "legend_handles = [\n",
    "    mpatches.Patch(color=\"#111111\", label='Klasa 0'),\n",
    "    mpatches.Patch(color=\"#bbbbbb\", label='Klasa 1')\n",
    "]\n",
    "plt.legend(handles=legend_handles, title='Klasy')\n",
    "plt.tight_layout()\n",
    "plt.savefig('rys_6.4.tiff', pil_kwargs={'compression': 'tiff_lzw'})\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
