{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wprowadzenie\n",
    "\n",
    "Niniejszy notatnik przedstawia zastosowanie grafowych sieci neuronowych do klasyfikacji stron na portalu Facebook. Wykorzystany zostanie zbiór danych zawierający informacje o wzajemnych powiązaniach między stronami oraz ich cechach charakterystycznych.\n",
    "\n",
    "Zadanie polega na przyporządkowaniu stron do odpowiednich kategorii tematycznych na podstawie:\n",
    "- struktury połączeń między stronami (założenie: strony o podobnej tematyce częściej odwołują się do siebie nawzajem)\n",
    "- cech numerycznych wprowadzonych przez administratorów stron\n",
    "\n",
    "Zbiór danych pochodzi z repozytorium SNAP (Stanford Network Analysis Project).\n",
    "\n",
    "Źródło: [zbiór danych SNAP dataset](https://snap.stanford.edu/data/facebook-large-page-page-network.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Załadowanie bibliotek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import zipfile\n",
    "from pathlib import Path\n",
    "from typing import Literal\n",
    "import os\n",
    "\n",
    "import networkx as nx\n",
    "import optuna\n",
    "import optuna.visualization as optvis\n",
    "import pandas as pd\n",
    "\n",
    "# Analiza najlepszych parametrów i modeli\n",
    "import plotly.express as px\n",
    "import pytorch_lightning as pl\n",
    "import requests\n",
    "import torch as th\n",
    "import torch.nn.functional as F\n",
    "import torch_geometric as pyg\n",
    "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n",
    "from torch_geometric.loader import NeighborLoader\n",
    "from torch_geometric.nn import ClusterGCNConv, GATv2Conv, GCNConv, SAGEConv\n",
    "from torch_geometric.utils import to_networkx\n",
    "from torchmetrics import AUROC, F1Score\n",
    "from IPython.display import display\n",
    "\n",
    "import sys\n",
    "sys.path.append('..')\n",
    "import utils.vis_utils as vis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import facebook_data_preprocessor as fdp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wczytanie danych i ich przetworzenie"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pobieranie zbioru danych Facebook...\n",
      "Rozpakowywanie zbioru danych...\n",
      "Wczytywanie plików danych...\n",
      "Przetwarzanie krawędzi...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/filip/projects/grafowe_sieci_neuronowe/rozdzial_6_zastosowania_sieci_gnn/klasyfikacja_wierzcholkow/facebook_data_preprocessor.py:155: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at /pytorch/torch/csrc/utils/tensor_new.cpp:253.)\n",
      "  edge_index = torch.tensor(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Przetwarzanie cech...\n",
      "Przetwarzanie etykiet docelowych...\n",
      "Tworzenie masek treningowych/walidacyjnych/testowych...\n"
     ]
    }
   ],
   "source": [
    "processor = fdp.FacebookDataPreprocessor(data_dir=\"./data\")\n",
    "processor.download_and_extract()\n",
    "data, le = processor.preprocess_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[22470, 31], edge_index=[2, 171002], y=[22470], train_mask=[22470], val_mask=[22470], test_mask=[22470])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_cls = th.unique(data.y)\n",
    "n_features = data.x.shape[1]\n",
    "n_train = data.train_mask.sum()\n",
    "n_val = data.val_mask.sum()\n",
    "n_test = data.test_mask.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "N train:  tensor(15730)\n",
      "N val:  tensor(3370)\n",
      "N test:  tensor(3370)\n"
     ]
    }
   ],
   "source": [
    "print(\"N train: \", n_train)\n",
    "print(\"N val: \", n_val)\n",
    "print(\"N test: \", n_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "government    0.306\n",
       "company       0.289\n",
       "politician    0.257\n",
       "tvshow        0.148\n",
       "Name: proportion, dtype: float64"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.Series(le.inverse_transform(data.y.numpy())).value_counts(normalize=True).round(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_undirected = pyg.transforms.ToUndirected()(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing METIS partitioning...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "cluster_data = pyg.loader.ClusterData(data_undirected, num_parts=15, recursive=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analiza danych grafowych"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Podstawowe właściwości grafu:\n",
      "- Liczba wierzchołków: 22470\n",
      "- Liczba krawędzi: 171002\n",
      "- Średni stopień wierzchołka: 15.22\n",
      "- Liczba składowych spójnych: 1\n",
      "- Gęstość grafu: 0.0007\n",
      "- Średni współczynnik klasteryzacji: 0.3597\n",
      "\n",
      "Analiza stopni wierzchołków:\n",
      "- Średni stopień: 15.22\n",
      "- Maksymalny stopień: 709\n",
      "- Minimalny stopień: 1\n",
      "\n",
      "Analiza największej składowej spójnej:\n",
      "- Rozmiar największej składowej: 22470\n",
      "- Procent wierzchołków w największej składowej: 100.00%\n"
     ]
    }
   ],
   "source": [
    "# Konwersja grafu do formatu networkx w celu analizy\n",
    "G = to_networkx(data, to_undirected=True)\n",
    "\n",
    "# Analiza podstawowych właściwości grafu\n",
    "n_nodes = G.number_of_nodes()\n",
    "n_edges = G.number_of_edges()\n",
    "avg_degree = 2 * n_edges / n_nodes\n",
    "n_components = nx.number_connected_components(G)\n",
    "density = nx.density(G)\n",
    "avg_clustering = nx.average_clustering(G)\n",
    "\n",
    "print(f\"Podstawowe właściwości grafu:\")\n",
    "print(f\"- Liczba wierzchołków: {n_nodes}\")\n",
    "print(f\"- Liczba krawędzi: {n_edges}\")\n",
    "print(f\"- Średni stopień wierzchołka: {avg_degree:.2f}\")\n",
    "print(f\"- Liczba składowych spójnych: {n_components}\")\n",
    "print(f\"- Gęstość grafu: {density:.4f}\")\n",
    "print(f\"- Średni współczynnik klasteryzacji: {avg_clustering:.4f}\")\n",
    "\n",
    "# Analiza rozkładu stopni wierzchołków\n",
    "degrees = [d for n, d in G.degree()]\n",
    "avg_degree = sum(degrees) / len(degrees)\n",
    "max_degree = max(degrees)\n",
    "min_degree = min(degrees)\n",
    "\n",
    "print(f\"\\nAnaliza stopni wierzchołków:\")\n",
    "print(f\"- Średni stopień: {avg_degree:.2f}\")\n",
    "print(f\"- Maksymalny stopień: {max_degree}\")\n",
    "print(f\"- Minimalny stopień: {min_degree}\")\n",
    "\n",
    "# Analiza największej składowej spójnej\n",
    "largest_cc = max(nx.connected_components(G), key=len)\n",
    "largest_cc_size = len(largest_cc)\n",
    "largest_cc_percentage = (largest_cc_size / n_nodes) * 100\n",
    "\n",
    "print(f\"\\nAnaliza największej składowej spójnej:\")\n",
    "print(f\"- Rozmiar największej składowej: {largest_cc_size}\")\n",
    "print(f\"- Procent wierzchołków w największej składowej: {largest_cc_percentage:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Analiza spójności grafu:\n",
      "- Czy graf jest spójny: Tak\n"
     ]
    }
   ],
   "source": [
    "# Sprawdzenie spójności grafu\n",
    "is_connected = nx.is_connected(G)\n",
    "\n",
    "print(\"\\nAnaliza spójności grafu:\")\n",
    "print(f\"- Czy graf jest spójny: {'Tak' if is_connected else 'Nie'}\")\n",
    "\n",
    "if not is_connected:\n",
    "    # Jeśli graf nie jest spójny, wyświetl dodatkowe informacje o składowych\n",
    "    components = list(nx.connected_components(G))\n",
    "    print(f\"- Liczba rozłącznych składowych: {len(components)}\")\n",
    "    print(f\"- Rozmiary składowych (od największej):\")\n",
    "    component_sizes = sorted([len(c) for c in components], reverse=True)\n",
    "    for i, size in enumerate(component_sizes[:5], 1):\n",
    "        print(f\"  {i}. składowa: {size} wierzchołków\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przygotowanie modelu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Definicja modelu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GNNModel(th.nn.Module):\n",
    "    \"\"\"\n",
    "    Model sieci neuronowej typu Graph Neural Network (GNN) do klasyfikacji wierzchołków w grafie.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    layer_size : int\n",
    "        Wielkość warstw ukrytych.\n",
    "    num_layers: int\n",
    "        Liczba warstw ukrytych.\n",
    "    n_cls : int\n",
    "        Liczba klas wyjściowych (liczba możliwych kategorii klasyfikacji).\n",
    "    conv_layer : torch_geometric.nn.conv\n",
    "        Typ warstwy konwolucyjnej grafu (domyślnie GCNConv). Możliwe opcje to GCNConv,\n",
    "        SAGEConv lub GATv2Conv.\n",
    "    **conv_kwargs : dict\n",
    "        Dodatkowe argumenty przekazywane do warstwy konwolucyjnej (np. liczba głów\n",
    "        w przypadku GATv2Conv).\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "    Model składa się z sekwencji warstw konwolucyjnych grafu, po których następuje\n",
    "    nieliniowość ReLU. Opcjonalnie, jeśli ostatnia warstwa nie ma rozmiaru\n",
    "    odpowiadającego liczbie klas, dodawana jest warstwa liniowa.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        layer_size: int,\n",
    "        num_layers: int,\n",
    "        n_features: int,\n",
    "        n_cls,\n",
    "        conv_layer=GCNConv,\n",
    "        batch_norm: bool = True,\n",
    "        jk: Literal[\"cat\", \"max\", \"lstm\"] = \"lstm\",\n",
    "        **conv_kwargs,\n",
    "    ):\n",
    "        super(GNNModel, self).__init__()\n",
    "        self.convs = th.nn.ModuleList()\n",
    "        self.dense = None\n",
    "        self.n_cls = n_cls\n",
    "        last_in = n_features\n",
    "        self.batch_norm = pyg.nn.BatchNorm(n_features) if batch_norm else None\n",
    "        # Iteracja po rozmiarach warstw zdefiniowanych w layer_sizes\n",
    "        for _ in range(num_layers):\n",
    "            # Dodanie warstwy konwolucyjnej do listy\n",
    "            self.convs.append(\n",
    "                conv_layer(in_channels=last_in, out_channels=layer_size, **conv_kwargs)\n",
    "            )\n",
    "            # W przypadku warstwy GATv2Conv z wieloma głowami uwagi (attention heads)\n",
    "            # i włączoną konkatenacją, wymiar wyjściowy zwiększa się\n",
    "            # Przykład: dla 4 głów i dim_out=32, faktyczny wymiar wyjściowy to 32*4=128\n",
    "            dim_out = layer_size\n",
    "            if conv_layer == GATv2Conv:\n",
    "                num_heads = conv_kwargs.get(\"heads\", 1)\n",
    "                concat = conv_kwargs.get(\"concat\", True)\n",
    "                if num_heads > 1 and concat:\n",
    "                    dim_out = layer_size * num_heads\n",
    "            last_in = dim_out\n",
    "        self.jk = (\n",
    "            pyg.nn.JumpingKnowledge(jk, channels=dim_out, num_layers=num_layers)\n",
    "            if jk\n",
    "            else None\n",
    "        )\n",
    "        # Jeśli ostatni wymiar nie odpowiada liczbie klas, dodaj warstwę liniową\n",
    "        if last_in != self.n_cls:\n",
    "            dense_in = last_in if jk != 'cat' else last_in * num_layers\n",
    "            self.dense = th.nn.Linear(dense_in, self.n_cls)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        \"\"\"\n",
    "        Funkcja służąca do przetwarzania danych wejściowych i generowania predykcji modelu.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : torch.Tensor\n",
    "            Macierz cech wierzchołków o wymiarach [liczba_wierzchołków x liczba_cech].\n",
    "        edge_index : torch.Tensor\n",
    "            Tensor krawędzi o wymiarach [2 x liczba_krawędzi].\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Przewidywane prawdopodobieństwa klas dla każdego wierzchołka.\n",
    "        \"\"\"\n",
    "        if self.batch_norm:\n",
    "            x = self.batch_norm(x)\n",
    "        output = []\n",
    "        for conv in self.convs:\n",
    "            x = conv(x, edge_index).relu()\n",
    "            output.append(x)\n",
    "        if self.jk:\n",
    "            x = self.jk(output)\n",
    "        if self.dense:\n",
    "            x = self.dense(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class GNNLightningModule(pl.LightningModule):\n",
    "    \"\"\"\n",
    "    Moduł PyTorch Lightning opakowujący model GNN, ułatwiający trening i ewaluację.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model : GNNModel\n",
    "        Instancja modelu GNN do treningu.\n",
    "    data : torch_geometric.data.Data\n",
    "        Obiekt zawierający dane grafu (cechy, krawędzie, etykiety).\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "    Moduł implementuje pełny cykl treningu z metrykami F1-Score i AUROC dla zadania\n",
    "    wieloklasowego. Wykorzystuje optymalizator Adam z domyślną stopą uczenia 0.01.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, model):\n",
    "        super(GNNLightningModule, self).__init__()\n",
    "        self.model = model\n",
    "        self.train_f1 = F1Score(num_classes=n_cls.size(0), task=\"multiclass\")\n",
    "        self.val_f1 = F1Score(num_classes=n_cls.size(0), task=\"multiclass\")\n",
    "        self.test_f1 = F1Score(num_classes=n_cls.size(0), task=\"multiclass\")\n",
    "        self.train_auc = AUROC(num_classes=n_cls.size(0), task=\"multiclass\")\n",
    "        self.val_auc = AUROC(num_classes=n_cls.size(0), task=\"multiclass\")\n",
    "        self.test_auc = AUROC(num_classes=n_cls.size(0), task=\"multiclass\")\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        \"\"\"\n",
    "        Funkcja służąca do wykonania predykcji modelu na podstawie danych wejściowych.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : torch.Tensor\n",
    "            Macierz cech wierzchołków.\n",
    "        edge_index : torch.Tensor\n",
    "            Tensor krawędzi.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Wyniki predykcji modelu.\n",
    "        \"\"\"\n",
    "        return self.model(x, edge_index)\n",
    "\n",
    "    def step(self, batch, mask, prefix):\n",
    "        \"\"\"\n",
    "        Funkcja realizująca pojedynczy krok (treningu/walidacji/testu).\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        batch : torch_geometric.data.Batch\n",
    "            Pakiet danych wejściowych.\n",
    "        mask : torch.Tensor\n",
    "            Maska określająca, które wierzchołki należą do danego zbioru.\n",
    "        prefix : str\n",
    "            Prefiks używany do logowania metryk ('train', 'val' lub 'test').\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Wartość funkcji straty.\n",
    "        \"\"\"\n",
    "        out = self(batch.x, batch.edge_index)\n",
    "        loss = F.cross_entropy(out[mask], batch.y[mask])\n",
    "        preds = out[mask]\n",
    "        labels = batch.y[mask]\n",
    "        self.log(f\"{prefix}_loss\", loss, batch_size=mask.sum())\n",
    "        self.log(\n",
    "            f\"{prefix}_f1\",\n",
    "            getattr(self, f\"{prefix}_f1\")(preds, labels),\n",
    "            batch_size=mask.sum(),\n",
    "        )\n",
    "        self.log(\n",
    "            f\"{prefix}_auc\",\n",
    "            getattr(self, f\"{prefix}_auc\")(preds, labels),\n",
    "            batch_size=mask.sum(),\n",
    "        )\n",
    "        return loss\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        \"\"\"\n",
    "        Funkcja realizująca krok treningu.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        batch : torch_geometric.data.Batch\n",
    "            Pakiet danych treningowych.\n",
    "        batch_idx : int\n",
    "            Indeks aktualnej partii danych.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Wartość funkcji straty dla kroku treningu.\n",
    "        \"\"\"\n",
    "        return self.step(batch, batch.train_mask, \"train\")\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        \"\"\"\n",
    "        Funkcja realizująca krok walidacji.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        batch : torch_geometric.data.Batch\n",
    "            Pakiet danych walidacyjnych.\n",
    "        batch_idx : int\n",
    "            Indeks aktualnej partii danych.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Wartość funkcji straty dla kroku walidacji.\n",
    "        \"\"\"\n",
    "        return self.step(batch, batch.val_mask, \"val\")\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        \"\"\"\n",
    "        Funkcja realizująca krok testowy.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        batch : torch_geometric.data.Batch\n",
    "            Pakiet danych testowych.\n",
    "        batch_idx : int\n",
    "            Indeks aktualnej partii danych.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Wartość funkcji straty dla kroku testowego.\n",
    "        \"\"\"\n",
    "        return self.step(batch, batch.test_mask, \"test\")\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        \"\"\"\n",
    "        Funkcja konfigurująca optymalizator dla modelu.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.optim.Optimizer\n",
    "            Skonfigurowany optymalizator Adam ze stopą uczenia 0.01.\n",
    "        \"\"\"\n",
    "        return th.optim.Adam(self.parameters(), lr=0.01)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_cluster_loader = pyg.loader.ClusterLoader(\n",
    "    cluster_data, batch_size=256, shuffle=False\n",
    ")\n",
    "\n",
    "val_cluster_loader = pyg.loader.ClusterLoader(\n",
    "    cluster_data, batch_size=256, shuffle=False\n",
    ")\n",
    "\n",
    "test_cluster_loader = pyg.loader.ClusterLoader(\n",
    "    cluster_data, batch_size=256, shuffle=False\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Strojenie hiperparametrów"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_layer_mapping = {\n",
    "    \"GCNConv\": GCNConv,\n",
    "    \"SAGEConv\": SAGEConv,\n",
    "    \"GATv2Conv\": GATv2Conv,\n",
    "    \"ClusterGCNConv\": ClusterGCNConv,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def enrich_model_params(params, n_features, n_cls):\n",
    "    \"\"\"\n",
    "    Wzbogaca słownik parametrów modelu o najlepsze hiperparametry oraz ogólne parametry.\n",
    "\n",
    "    Parametry\n",
    "    ----------\n",
    "    params : dict\n",
    "        Słownik zawierający najlepsze hiperparametry.\n",
    "    n_features : int\n",
    "        Liczba cech.\n",
    "    n_cls : int\n",
    "        Liczba klas.\n",
    "\n",
    "    Zwraca\n",
    "    -------\n",
    "    dict\n",
    "        Wzbogacony słownik parametrów modelu.\n",
    "    \"\"\"\n",
    "    model_params = {\n",
    "        'n_features': n_features,\n",
    "        'n_cls': n_cls.shape[0],\n",
    "        'conv_layer': conv_layer_mapping[params[\"conv_layer\"]],\n",
    "        'num_layers': params[\"num_layers\"],\n",
    "        'layer_size': params[\"layer_size\"],\n",
    "        'jk': params.get(\"jk\", None),\n",
    "        'batch_norm': params.get(\"batch_norm\", True),\n",
    "    }\n",
    "\n",
    "    if params[\"conv_layer\"] == \"GATv2Conv\":\n",
    "        model_params[\"dropout\"] = params.get(\"dropout\", 0)\n",
    "        model_params[\"heads\"] = params.get(\"num_heads\", 1)\n",
    "    elif params[\"conv_layer\"] == \"ClusterGCNConv\":\n",
    "        model_params[\"project\"] = params.get(\"project\", False)\n",
    "    elif params[\"conv_layer\"] == \"SAGEConv\":\n",
    "        model_params[\"normalize\"] = params.get(\"normalize\", False)\n",
    "        model_params[\"project\"] = params.get(\"project\", False)\n",
    "    elif params[\"conv_layer\"] == \"GCNConv\":\n",
    "        model_params[\"normalize\"] = params.get(\"normalize\", False)\n",
    "\n",
    "    return model_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "conv_layer_name_2_class = {v: k for k, v in conv_layer_mapping.items()}\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",
    "    # Define hyperparameters to tune\n",
    "    conv_layer_name = trial.suggest_categorical(\n",
    "        \"conv_layer\", list(conv_layer_mapping.keys())\n",
    "    )\n",
    "    conv_layer = conv_layer_mapping[conv_layer_name]\n",
    "\n",
    "    batch_norm = trial.suggest_categorical(\"batch_norm\", [True, False])\n",
    "    if conv_layer_name != \"SAGEConv\":\n",
    "        aggr = trial.suggest_categorical(\"aggr_single\", [\"sum\", \"mean\", \"max\"])\n",
    "    else:\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",
    "\n",
    "    num_layers = trial.suggest_int(\"num_layers\", 2, 5)\n",
    "    layer_size = trial.suggest_categorical(\"layer_size\", [8, 16, 32])\n",
    "    jk = trial.suggest_categorical(\"jk\", [\"cat\", \"max\", \"lstm\", None])\n",
    "    # Przygotowanie modelu\n",
    "    pyg.seed_everything(123)\n",
    "    if conv_layer_name == \"GATv2Conv\":\n",
    "        num_heads = trial.suggest_int(\"num_heads\", 1, 3)\n",
    "        dropout = trial.suggest_categorical(\"dropout\", [0, 0.2, 0.5])\n",
    "        gnn = GNNModel(\n",
    "            num_layers=num_layers,\n",
    "            layer_size=layer_size,\n",
    "            n_features=n_features,\n",
    "            n_cls=n_cls.shape[0],\n",
    "            conv_layer=conv_layer,\n",
    "            aggr=aggr,\n",
    "            heads=num_heads,\n",
    "            batch_norm=batch_norm,\n",
    "            dropout=dropout,\n",
    "            jk=jk,\n",
    "        )\n",
    "    elif conv_layer_name == \"SAGEConv\":\n",
    "        project = trial.suggest_categorical(\"project\", [True, False])\n",
    "        normalize = trial.suggest_categorical(\"normalize\", [True, False])\n",
    "        gnn = GNNModel(\n",
    "            num_layers=num_layers,\n",
    "            layer_size=layer_size,\n",
    "            n_features=n_features,\n",
    "            n_cls=n_cls.shape[0],\n",
    "            conv_layer=conv_layer,\n",
    "            aggr=aggr,\n",
    "            batch_norm=batch_norm,\n",
    "            project=project,\n",
    "            normalize=normalize,\n",
    "            jk=jk,\n",
    "        )\n",
    "    elif conv_layer_name == \"GCNConv\":\n",
    "        normalize = trial.suggest_categorical(\"normalize\", [True, False])\n",
    "        gnn = GNNModel(\n",
    "            num_layers=num_layers,\n",
    "            layer_size=layer_size,\n",
    "            n_features=n_features,\n",
    "            n_cls=n_cls.shape[0],\n",
    "            conv_layer=conv_layer,\n",
    "            aggr=aggr,\n",
    "            batch_norm=batch_norm,\n",
    "            normalize=normalize,\n",
    "            jk=jk,\n",
    "        )\n",
    "    else:\n",
    "        gnn = GNNModel(\n",
    "            num_layers=num_layers,\n",
    "            layer_size=layer_size,\n",
    "            n_features=n_features,\n",
    "            n_cls=n_cls.shape[0],\n",
    "            conv_layer=conv_layer,\n",
    "            aggr=aggr,\n",
    "            batch_norm=batch_norm,\n",
    "            jk=jk,\n",
    "        )\n",
    "\n",
    "    model = GNNLightningModule(gnn)\n",
    "\n",
    "    # Tworzenie trenera sieci\n",
    "    trainer = pl.Trainer(\n",
    "        max_epochs=50,\n",
    "        enable_progress_bar=False,\n",
    "        callbacks=[EarlyStopping(monitor=\"val_f1\", mode=\"max\", patience=5)],\n",
    "        enable_checkpointing=False,\n",
    "        logger=False,\n",
    "        deterministic=True,\n",
    "    )\n",
    "\n",
    "    # Trening i walidacja modelu\n",
    "    trainer.fit(model, train_cluster_loader, val_cluster_loader)\n",
    "\n",
    "    # Pobranie najlepszego wyniku F1 na zbiorze walidacyjnym\n",
    "    val_f1 = trainer.callback_metrics[\"val_f1\"].item()\n",
    "    return val_f1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wczytane najlepsze hiperparametry:  {'n_features': 31, 'n_cls': 4, 'conv_layer': <class 'torch_geometric.nn.conv.sage_conv.SAGEConv'>, 'num_layers': 3, 'layer_size': 32, 'jk': 'cat', 'batch_norm': True, 'normalize': True, 'project': True}\n"
     ]
    }
   ],
   "source": [
    "# Przełącznik czy stroić parametry, czy wczytać z pliku\n",
    "tune_params = False\n",
    "\n",
    "if tune_params:\n",
    "    # Tworzenie eksperymentu Optuna i optymalizacja funkcji celu\n",
    "    study = optuna.create_study(direction=\"maximize\")\n",
    "    study.optimize(objective, n_trials=50)\n",
    "\n",
    "    # Wyświetlenie najlepszych hiperparametrów\n",
    "    best_params = study.best_params\n",
    "    print(\"Najlepsze hiperparametry: \", best_params)\n",
    "\n",
    "    # Zapisanie najlepszych hiperparametrów do pliku\n",
    "    with open(\"facebook_best.json\", \"w\") as f:\n",
    "        json.dump(best_params, f)\n",
    "    best_params = enrich_model_params(best_params, n_features, n_cls)\n",
    "else:\n",
    "    # Wczytanie najlepszych hiperparametrów z pliku\n",
    "    with open(\"facebook_best.json\", \"r\") as f:\n",
    "        best_params = json.load(f)\n",
    "    best_params = enrich_model_params(best_params, n_features, n_cls)\n",
    "    print(\"Wczytane najlepsze hiperparametry: \", best_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sprawdzenie wyników strojenia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "if tune_params:\n",
    "    display(study.trials_dataframe().sort_values(\"value\", ascending=False).head(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "if tune_params:\n",
    "    # Pobranie ramki danych z wynikami\n",
    "    df = study.trials_dataframe()\n",
    "\n",
    "    # Wizualizacja ważności parametrów\n",
    "    fig = optvis.plot_param_importances(study)\n",
    "    fig.show()\n",
    "\n",
    "    # Szczegółowa analiza najlepszych modeli\n",
    "    print(\"\\nSzczegółowa analiza 5 najlepszych modeli:\")\n",
    "    top_5_trials = df.nlargest(5, \"value\")\n",
    "    for idx, row in top_5_trials.iterrows():\n",
    "        print(f\"\\nModel {idx + 1}:\")\n",
    "        print(f\"Wartość F1: {row['value']:.4f}\")\n",
    "        print(\"Parametry:\")\n",
    "        params = {\n",
    "            k.replace(\"params_\", \"\"): v for k, v in row.items() if k.startswith(\"params_\")\n",
    "        }\n",
    "        for param, value in params.items():\n",
    "            if not isinstance(value, list) and not pd.isna(value):\n",
    "                print(f\"- {param}: {value}\")\n",
    "\n",
    "    # Analiza rozkładu wartości dla poszczególnych parametrów\n",
    "    param_names = [name for name in df.columns if name.startswith(\"params_\")]\n",
    "    # Korelacja między parametrami a wynikiem\n",
    "    correlations = {}\n",
    "    for param in param_names:\n",
    "        if df[param].dtype in [\"int64\", \"float64\"]:\n",
    "            correlations[param.replace(\"params_\", \"\")] = df[param].corr(df[\"value\"])\n",
    "\n",
    "    print(\"\\nKorelacja parametrów z wynikiem F1:\")\n",
    "    for param, corr in sorted(correlations.items(), key=lambda x: abs(x[1]), reverse=True):\n",
    "        print(f\"{param}: {corr:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Szkolenie modelu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Najlepsze parametry:\n",
      "n_features: 31\n",
      "n_cls: 4\n",
      "conv_layer: <class 'torch_geometric.nn.conv.sage_conv.SAGEConv'>\n",
      "num_layers: 3\n",
      "layer_size: 32\n",
      "jk: cat\n",
      "batch_norm: True\n",
      "normalize: True\n",
      "project: True\n"
     ]
    }
   ],
   "source": [
    "print(\"Najlepsze parametry:\")\n",
    "for key, value in best_params.items():\n",
    "    print(f\"{key}: {value}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "if th.cuda.is_available():\n",
    "    # Dodatkowe ustawienie determinizmu dla GPU\n",
    "    os.environ[\"CUBLAS_WORKSPACE_CONFIG\"]=\":16:8\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inicjalizacja modelu z najlepszymi parametrami\n",
    "pyg.seed_everything(123)\n",
    "model = GNNModel(**best_params)\n",
    "\n",
    "model_pl = GNNLightningModule(model)\n",
    "\n",
    "# Inicjalizacja trenera\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=\"best_model\",\n",
    ")\n",
    "\n",
    "trainer = pl.Trainer(\n",
    "    max_epochs=100,\n",
    "    accelerator=\"cpu\",\n",
    "    logger=False,\n",
    "    enable_model_summary=True,\n",
    "    deterministic=True,\n",
    "    callbacks=[\n",
    "        # EarlyStopping(monitor=\"val_f1\", mode=\"max\", patience=5),\n",
    "        checkpoint_callback,\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Sprawdzenie metryk przed rozpoczęciem treningu\n",
    "val_scores_before = trainer.validate(model_pl, val_cluster_loader)\n",
    "\n",
    "# Trenowanie modelu\n",
    "trainer.fit(model_pl, train_cluster_loader, val_cluster_loader)\n",
    "\n",
    "# Wczytanie najlepszego modelu\n",
    "best_model = GNNLightningModule.load_from_checkpoint(\n",
    "    checkpoint_callback.best_model_path, model=model\n",
    ")\n",
    "\n",
    "# Ewaluacja najlepszego modelu\n",
    "val_scores_after = trainer.validate(best_model, val_cluster_loader)\n",
    "test_results = trainer.test(best_model, test_cluster_loader)\n",
    "\n",
    "print(\"\\nWyniki:\")\n",
    "print(\"Przed treningiem:\")\n",
    "print(\n",
    "    f\"Zbiór walidacyjny - F1: {val_scores_before[0]['val_f1']:.4f}, AUC: {val_scores_before[0]['val_auc']:.4f}\"\n",
    ")\n",
    "print(\"Po treningu:\")\n",
    "print(\n",
    "    f\"Zbiór walidacyjny - F1: {val_scores_after[0]['val_f1']:.4f}, AUC: {val_scores_after[0]['val_auc']:.4f}\"\n",
    ")\n",
    "print(\n",
    "    f\"Zbiór testowy - F1: {test_results[0]['test_f1']:.4f}, AUC: {test_results[0]['test_auc']:.4f}\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rzutowanie reprezentacji wektorowej na płaszczyznę"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 639164 characters]",
      "text/plain": [
       "<Figure size 787.402x787.402 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "from sklearn.manifold import TSNE\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Funkcja do redukcji wymiarowości i wizualizacji\n",
    "def visualize_embeddings(model, data, le, method='PCA'):\n",
    "    \"\"\"\n",
    "    Funkcja do wizualizacji reprezentacji wierzchołków na płaszczyźnie za pomocą PCA lub T-SNE.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model : GNNLightningModule\n",
    "        Wytrenowany model GNN.\n",
    "    data : torch_geometric.data.Data\n",
    "        Dane grafowe zawierające cechy wierzchołków, krawędzie i etykiety.\n",
    "    le : LabelEncoder\n",
    "        Label encoder używany do transformacji etykiet.\n",
    "    method : str\n",
    "        Metoda redukcji wymiarowości ('PCA' lub 'T-SNE').\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    None\n",
    "    \"\"\"\n",
    "    # Przetwarzanie danych przez model w celu uzyskania reprezentacji wierzchołków\n",
    "    model.eval()\n",
    "    with th.no_grad():\n",
    "        embeddings = model(data.x, data.edge_index).cpu().numpy()\n",
    "\n",
    "    # Redukcja wymiarowości\n",
    "    if method == 'PCA':\n",
    "        reducer = PCA(n_components=2)\n",
    "    elif method == 'T-SNE':\n",
    "        reducer = TSNE(n_components=2, random_state=42)\n",
    "    else:\n",
    "        raise ValueError(\"Metoda redukcji wymiarowości musi być 'PCA' lub 'T-SNE'\")\n",
    "\n",
    "    reduced_embeddings = reducer.fit_transform(embeddings)\n",
    "\n",
    "    # Wizualizacja\n",
    "    plt.rcParams.update({'font.size': 12})\n",
    "    fig, ax = plt.subplots(figsize=vis.image_size_in_cm(20, 20))\n",
    "    \n",
    "    # Użycie niestandardowej palety kolorów w odcieniach szarości z lepszym kontrastem\n",
    "    colors = ['#000000', '#404040', '#808080', '#C0C0C0']  # Czarny, ciemnoszary, średnioszary, jasnoszary\n",
    "    scatter = ax.scatter(reduced_embeddings[:, 0], reduced_embeddings[:, 1], c=[colors[i] for i in data.y], alpha=0.7)\n",
    "\n",
    "    unique_classes = th.unique(data.y).numpy()\n",
    "    handles = [plt.Line2D([0], [0], marker='o', color='w', markerfacecolor=colors[i], markersize=5) for i in unique_classes]\n",
    "    labels = le.inverse_transform(unique_classes)\n",
    "    vis.change_plot_comma_separator(ax, x_change=True, y_change=True)\n",
    "    plt.legend(handles, labels, title='Klasa')\n",
    "    plt.xlabel('Wymiar 1')\n",
    "    plt.ylabel('Wymiar 2')\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(\"rys_6.1.tiff\", format='tiff', pil_kwargs={'compression': 'tiff_lzw'})\n",
    "    plt.show()\n",
    "\n",
    "# Wywołanie funkcji wizualizującej z wybraną metodą redukcji wymiarowości\n",
    "visualize_embeddings(best_model, data, le, method='PCA')  # Można zmienić na 'T-SNE'"
   ]
  }
 ],
 "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
}
