graph_ecg_v1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "787797a9-7447-47a6-9cb1-20564b0b0d12",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from sklearn.metrics import accuracy_score, confusion_matrix, classification_report, roc_curve, auc, precision_recall_curve\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.loader import DataLoader\n",
    "from torch_geometric.nn import GATConv, global_mean_pool, global_add_pool, global_max_pool\n",
    "from torch_geometric.utils import add_self_loops, dense_to_sparse\n",
    "from torch.utils.data import TensorDataset\n",
    "from tslearn.datasets import UCR_UEA_datasets\n",
    "import warnings\n",
    "from scipy import stats\n",
    "from collections import defaultdict\n",
    "import pandas as pd\n",
    "from scipy.signal import find_peaks\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# Configuration\n",
    "RANDOM_SEED = 42\n",
    "np.random.seed(RANDOM_SEED)\n",
    "torch.manual_seed(RANDOM_SEED)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed(RANDOM_SEED)\n",
    "\n",
    "class Config:\n",
    "    MAX_EPOCHS = 200\n",
    "    BATCH_SIZE = 32\n",
    "    LR = 0.001\n",
    "    N_SPLITS = 5\n",
    "    HIDDEN_DIM = 128\n",
    "    HEADS = 8\n",
    "    DROPOUT = 0.3\n",
    "    WEIGHT_DECAY = 1e-4\n",
    "    PATIENCE = 15\n",
    "    PLOT_DIR = \"./plots\"\n",
    "    RESULTS_DIR = \"./results\"\n",
    "    \n",
    "    # Model architecture\n",
    "    NUM_GAT_LAYERS = 3\n",
    "    USE_EDGE_FEATURES = True\n",
    "    POOLING_STRATEGY = 'hybrid'  # 'mean', 'max', 'add', 'hybrid'\n",
    "    USE_RESIDUAL = True\n",
    "    USE_BATCH_NORM = True\n",
    "    \n",
    "    # Novel components\n",
    "    USE_TEMPORAL_ENCODING = True\n",
    "    USE_MULTI_SCALE = True\n",
    "    USE_WAVELET_FEATURES = True\n",
    "    \n",
    "    def __init__(self):\n",
    "        os.makedirs(self.PLOT_DIR, exist_ok=True)\n",
    "        os.makedirs(self.RESULTS_DIR, exist_ok=True)\n",
    "\n",
    "config = Config()\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "def analyze_dataset(X, y, dataset_name):\n",
    "    \"\"\"Comprehensive dataset analysis and visualization.\"\"\"\n",
    "    stats_dict = {\n",
    "        'Dataset Name': dataset_name,\n",
    "        'Number of Samples': len(X),\n",
    "        'Time Series Length': X.shape[1],\n",
    "        'Number of Classes': len(np.unique(y)),\n",
    "        'Class Distribution': dict(zip(*np.unique(y, return_counts=True))),\n",
    "        'Mean Length': X.shape[1],\n",
    "        'Mean Value': X.mean(),\n",
    "        'Std Value': X.std(),\n",
    "        'Max Value': X.max(),\n",
    "        'Min Value': X.min()\n",
    "    }\n",
    "    \n",
    "    # Create visualization\n",
    "    fig = plt.figure(figsize=(15, 10))\n",
    "    gs = plt.GridSpec(3, 2)\n",
    "    \n",
    "    # Plot 1: Example time series from each class\n",
    "    ax1 = fig.add_subplot(gs[0, :])\n",
    "    classes = np.unique(y)\n",
    "    colors = plt.cm.tab10(np.linspace(0, 1, len(classes)))\n",
    "    for i, cls in enumerate(classes):\n",
    "        idx = np.where(y == cls)[0][0]\n",
    "        ax1.plot(X[idx], color=colors[i], label=f'Class {cls}', alpha=0.8)\n",
    "    ax1.set_title('Example Time Series from Each Class')\n",
    "    ax1.legend()\n",
    "    ax1.grid(True, alpha=0.3)\n",
    "    \n",
    "    # Plot 2: Class distribution\n",
    "    ax2 = fig.add_subplot(gs[1, 0])\n",
    "    sns.barplot(x=list(stats_dict['Class Distribution'].keys()),\n",
    "                y=list(stats_dict['Class Distribution'].values()),\n",
    "                palette='tab10', ax=ax2)\n",
    "    ax2.set_title('Class Distribution')\n",
    "    \n",
    "    # Plot 3: Signal characteristics\n",
    "    ax3 = fig.add_subplot(gs[1, 1])\n",
    "    for i, cls in enumerate(classes):\n",
    "        class_data = X[y == cls].flatten()\n",
    "        sns.kdeplot(data=class_data, ax=ax3, label=f'Class {cls}')\n",
    "    ax3.set_title('Signal Distribution by Class')\n",
    "    ax3.legend()\n",
    "    \n",
    "    # Plot 4: Peaks and valleys analysis\n",
    "    ax4 = fig.add_subplot(gs[2, :])\n",
    "    example_idx = np.where(y == classes[0])[0][0]\n",
    "    example_series = X[example_idx]\n",
    "    peaks, _ = find_peaks(example_series)\n",
    "    valleys, _ = find_peaks(-example_series)\n",
    "    \n",
    "    ax4.plot(example_series, label='Signal')\n",
    "    ax4.plot(peaks, example_series[peaks], \"x\", label='Peaks')\n",
    "    ax4.plot(valleys, example_series[valleys], \"o\", label='Valleys')\n",
    "    ax4.set_title('Peak and Valley Analysis')\n",
    "    ax4.legend()\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.savefig(os.path.join(config.PLOT_DIR, f'{dataset_name}_analysis.png'), dpi=300, bbox_inches='tight')\n",
    "    plt.close()\n",
    "    \n",
    "    return pd.DataFrame([stats_dict]).T\n",
    "\n",
    "def extract_wavelets(X):\n",
    "    \"\"\"Extract wavelet features from time series.\"\"\"\n",
    "    import pywt\n",
    "    wavelets = []\n",
    "    for x in X:\n",
    "        coeffs = pywt.wavedec(x, 'db4', level=3)\n",
    "        features = np.concatenate([c for c in coeffs])\n",
    "        wavelets.append(features)\n",
    "    return np.array(wavelets)\n",
    "\n",
    "def compute_temporal_encoding(length, max_freq=10000):\n",
    "    \"\"\"Compute temporal position encoding.\"\"\"\n",
    "    position = np.arange(length)\n",
    "    freq_term = np.exp(-np.log(max_freq) * np.arange(64) / 64)\n",
    "    pos_enc = np.zeros((length, 64))\n",
    "    \n",
    "    for i in range(0, 64, 2):\n",
    "        pos_enc[:, i] = np.sin(position / freq_term[i])\n",
    "        pos_enc[:, i+1] = np.cos(position / freq_term[i])\n",
    "        \n",
    "    return torch.FloatTensor(pos_enc)\n",
    "\n",
    "class MultiScaleGATConv(nn.Module):\n",
    "    \"\"\"Multi-scale GAT layer with multiple window sizes.\"\"\"\n",
    "    def __init__(self, in_channels, out_channels, heads=8, window_sizes=[3, 5, 7]):\n",
    "        super().__init__()\n",
    "        self.window_sizes = window_sizes\n",
    "        self.convs = nn.ModuleList([\n",
    "            GATConv(in_channels, out_channels // len(window_sizes), \n",
    "                   heads=heads, dropout=config.DROPOUT)\n",
    "            for _ in window_sizes\n",
    "        ])\n",
    "        self.attention_weights = None\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        outputs = []\n",
    "        self.attention_weights = []\n",
    "        \n",
    "        for conv, window_size in zip(self.convs, self.window_sizes):\n",
    "            # Create edges for current window size\n",
    "            edge_list = []\n",
    "            num_nodes = x.size(0)\n",
    "            for i in range(num_nodes):\n",
    "                for j in range(max(0, i-window_size), min(num_nodes, i+window_size+1)):\n",
    "                    if i != j:\n",
    "                        edge_list.append([i, j])\n",
    "            \n",
    "            if edge_list:\n",
    "                window_edge_index = torch.tensor(edge_list, device=x.device).t()\n",
    "                window_edge_index = add_self_loops(window_edge_index, num_nodes=num_nodes)[0]\n",
    "                \n",
    "                # Get output and attention weights\n",
    "                out, (_, att_weights) = conv(x, window_edge_index, return_attention_weights=True)\n",
    "                outputs.append(out)\n",
    "                self.attention_weights.append((window_edge_index, att_weights))\n",
    "            \n",
    "        return torch.cat(outputs, dim=-1)\n",
    "\n",
    "class EdgeFeatureGATConv(GATConv):\n",
    "    \"\"\"GAT layer with edge feature support.\"\"\"\n",
    "    def __init__(self, in_channels, out_channels, heads=8, edge_dim=1):\n",
    "        super().__init__(in_channels, out_channels, heads=heads, edge_dim=edge_dim)\n",
    "        \n",
    "    def forward(self, x, edge_index, edge_attr=None):\n",
    "        if edge_attr is None:\n",
    "            # Create simple edge features based on node distance\n",
    "            edge_attr = torch.abs(edge_index[0] - edge_index[1]).float().unsqueeze(-1)\n",
    "        return super().forward(x, edge_index, edge_attr=edge_attr)\n",
    "\n",
    "class DynamicTemporalGraph(nn.Module):\n",
    "    \"\"\"Dynamic graph construction with learnable adjacency.\"\"\"\n",
    "    def __init__(self, input_dim, hidden_dim):\n",
    "        super().__init__()\n",
    "        self.feature_transform = nn.Sequential(\n",
    "            nn.Linear(input_dim, hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_dim, hidden_dim)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # Transform features\n",
    "        h = self.feature_transform(x)\n",
    "        \n",
    "        # Compute similarity matrix\n",
    "        sim_matrix = torch.matmul(h, h.transpose(-2, -1))\n",
    "        \n",
    "        # Apply softmax to get attention weights\n",
    "        adj_matrix = torch.softmax(sim_matrix / np.sqrt(h.size(-1)), dim=-1)\n",
    "        \n",
    "        # Convert to sparse format\n",
    "        edge_index, edge_attr = dense_to_sparse(adj_matrix)\n",
    "        \n",
    "        return edge_index, edge_attr\n",
    "\n",
    "class EnhancedGNNClassifier(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim, num_classes):\n",
    "        super().__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        \n",
    "        # Temporal position encoding\n",
    "        if config.USE_TEMPORAL_ENCODING:\n",
    "            self.temporal_enc = nn.Parameter(torch.randn(1000, input_dim))\n",
    "            input_dim = input_dim * 2\n",
    "            \n",
    "        # Feature transformation\n",
    "        self.feature_transform = nn.Linear(input_dim, hidden_dim)\n",
    "        \n",
    "        # Dynamic graph construction\n",
    "        self.graph_constructor = DynamicTemporalGraph(hidden_dim, hidden_dim)\n",
    "        \n",
    "        # Multi-scale GAT layers\n",
    "        self.gat_layers = nn.ModuleList([\n",
    "            MultiScaleGATConv(\n",
    "                hidden_dim if i == 0 else hidden_dim * config.HEADS,\n",
    "                hidden_dim,\n",
    "                heads=config.HEADS\n",
    "            )\n",
    "            for i in range(config.NUM_GAT_LAYERS)\n",
    "        ])\n",
    "        \n",
    "        # Batch normalization layers\n",
    "        if config.USE_BATCH_NORM:\n",
    "            self.batch_norms = nn.ModuleList([\n",
    "                nn.BatchNorm1d(hidden_dim * config.HEADS)\n",
    "                for _ in range(config.NUM_GAT_LAYERS)\n",
    "            ])\n",
    "        \n",
    "        # Pooling layers\n",
    "        if config.POOLING_STRATEGY == 'hybrid':\n",
    "            self.pool_dim = hidden_dim * config.HEADS * 3\n",
    "        else:\n",
    "            self.pool_dim = hidden_dim * config.HEADS\n",
    "            \n",
    "        # Classification head\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(self.pool_dim, hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(config.DROPOUT),\n",
    "            nn.Linear(hidden_dim, num_classes)\n",
    "        )\n",
    "        \n",
    "        # Store attention weights for interpretability\n",
    "        self.attention_weights = []\n",
    "        \n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "        self.attention_weights = []\n",
    "        \n",
    "        # Add temporal encoding if enabled\n",
    "        if config.USE_TEMPORAL_ENCODING:\n",
    "            seq_len = x.size(0)\n",
    "            x = torch.cat([x, self.temporal_enc[:seq_len]], dim=-1)\n",
    "        \n",
    "        # Initial feature transformation\n",
    "        x = self.feature_transform(x)\n",
    "        \n",
    "        # Dynamic graph construction\n",
    "        if config.USE_EDGE_FEATURES:\n",
    "            edge_index, edge_attr = self.graph_constructor(x)\n",
    "        else:\n",
    "            edge_attr = None\n",
    "            \n",
    "        # Apply GAT layers with residual connections\n",
    "        for i, gat_layer in enumerate(self.gat_layers):\n",
    "            identity = x\n",
    "            x = gat_layer(x, edge_index)\n",
    "            \n",
    "            if config.USE_BATCH_NORM:\n",
    "                x = self.batch_norms[i](x)\n",
    "                \n",
    "            x = torch.relu(x)\n",
    "            \n",
    "            if config.USE_RESIDUAL and x.size(-1) == identity.size(-1):\n",
    "                x = x + identity\n",
    "                \n",
    "            self.attention_weights.append(gat_layer.attention_weights)\n",
    "            \n",
    "        # Pooling\n",
    "        if config.POOLING_STRATEGY == 'hybrid':\n",
    "            x_mean = global_mean_pool(x, batch)\n",
    "            x_max = global_max_pool(x, batch)\n",
    "            x_add = global_add_pool(x, batch)\n",
    "            x = torch.cat([x_mean, x_max, x_add], dim=-1)\n",
    "        elif config.POOLING_STRATEGY == 'mean':\n",
    "            x = global_mean_pool(x, batch)\n",
    "        elif config.POOLING_STRATEGY == 'max':\n",
    "            x = global_max_pool(x, batch)\n",
    "            \n",
    "        # Classification\n",
    "        out = self.classifier(x)\n",
    "        return out\n",
    "    \n",
    "    def get_attention_maps(self, data):\n",
    "        \"\"\"Get attention maps for all layers.\"\"\"\n",
    "        self.eval()\n",
    "        with torch.no_grad():\n",
    "            _ = self.forward(data)\n",
    "            \n",
    "        attention_maps = []\n",
    "        for layer_weights in self.attention_weights:\n",
    "            layer_maps = []\n",
    "            for edge_index, weights in layer_weights:\n",
    "                # Convert to dense attention matrix\n",
    "                num_nodes = data.x.size(0)\n",
    "                dense_att = torch.zeros(num_nodes, num_nodes, device=data.x.device)\n",
    "                dense_att[edge_index[0], edge_index[1]] = weights.mean(dim=1)\n",
    "                layer_maps.append(dense_att.cpu().numpy())\n",
    "            attention_maps.append(layer_maps)\n",
    "            \n",
    "        return attention_maps\n",
    "    \n",
    "class TrainingManager:\n",
    "    def __init__(self, model, optimizer, criterion, device, config):\n",
    "        self.model = model\n",
    "        self.optimizer = optimizer\n",
    "        self.criterion = criterion\n",
    "        self.device = device\n",
    "        self.config = config\n",
    "        self.early_stopping = EarlyStopping(patience=config.PATIENCE)\n",
    "        self.history = defaultdict(list)\n",
    "        \n",
    "    def train_epoch(self, train_loader):\n",
    "        self.model.train()\n",
    "        total_loss = 0\n",
    "        predictions = []\n",
    "        targets = []\n",
    "        outputs = []\n",
    "        \n",
    "        for batch in train_loader:\n",
    "            batch = batch.to(self.device)\n",
    "            self.optimizer.zero_grad()\n",
    "            \n",
    "            output = self.model(batch)\n",
    "            loss = self.criterion(output, batch.y.squeeze())\n",
    "            \n",
    "            loss.backward()\n",
    "            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)\n",
    "            self.optimizer.step()\n",
    "            \n",
    "            total_loss += loss.item()\n",
    "            predictions.extend(output.argmax(dim=1).cpu().numpy())\n",
    "            targets.extend(batch.y.cpu().numpy())\n",
    "            outputs.extend(torch.softmax(output, dim=1).cpu().numpy())\n",
    "            \n",
    "        metrics = compute_metrics(targets, predictions, outputs)\n",
    "        metrics['loss'] = total_loss / len(train_loader)\n",
    "        \n",
    "        return metrics\n",
    "    \n",
    "    def evaluate(self, val_loader):\n",
    "        self.model.eval()\n",
    "        total_loss = 0\n",
    "        predictions = []\n",
    "        targets = []\n",
    "        outputs = []\n",
    "        attention_maps = []\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            for batch in val_loader:\n",
    "                batch = batch.to(self.device)\n",
    "                output = self.model(batch)\n",
    "                loss = self.criterion(output, batch.y.squeeze())\n",
    "                \n",
    "                total_loss += loss.item()\n",
    "                predictions.extend(output.argmax(dim=1).cpu().numpy())\n",
    "                targets.extend(batch.y.cpu().numpy())\n",
    "                outputs.extend(torch.softmax(output, dim=1).cpu().numpy())\n",
    "                \n",
    "                # Collect attention maps for first batch\n",
    "                if len(attention_maps) == 0:\n",
    "                    attention_maps = self.model.get_attention_maps(batch)\n",
    "        \n",
    "        metrics = compute_metrics(targets, predictions, outputs)\n",
    "        metrics['loss'] = total_loss / len(val_loader)\n",
    "        \n",
    "        return metrics, attention_maps\n",
    "    \n",
    "    def train(self, train_loader, val_loader, epochs):\n",
    "        best_val_auc = 0\n",
    "        best_model = None\n",
    "        \n",
    "        for epoch in range(epochs):\n",
    "            # Training phase\n",
    "            train_metrics = self.train_epoch(train_loader)\n",
    "            \n",
    "            # Validation phase\n",
    "            val_metrics, attention_maps = self.evaluate(val_loader)\n",
    "            \n",
    "            # Store history\n",
    "            for k, v in train_metrics.items():\n",
    "                self.history[f'train_{k}'].append(v)\n",
    "            for k, v in val_metrics.items():\n",
    "                self.history[f'val_{k}'].append(v)\n",
    "            \n",
    "            # Print progress\n",
    "            if epoch % 10 == 0:\n",
    "                print(f'Epoch {epoch}:')\n",
    "                print(f\"Train - Loss: {train_metrics['loss']:.4f}, \"\n",
    "                      f\"AUC: {train_metrics['auc']:.4f}, \"\n",
    "                      f\"Acc: {train_metrics['accuracy']:.4f}\")\n",
    "                print(f\"Val - Loss: {val_metrics['loss']:.4f}, \"\n",
    "                      f\"AUC: {val_metrics['auc']:.4f}, \"\n",
    "                      f\"Acc: {val_metrics['accuracy']:.4f}\")\n",
    "            \n",
    "            # Save best model\n",
    "            if val_metrics['auc'] > best_val_auc:\n",
    "                best_val_auc = val_metrics['auc']\n",
    "                best_model = copy.deepcopy(self.model.state_dict())\n",
    "                best_attention_maps = attention_maps\n",
    "            \n",
    "            # Early stopping\n",
    "            self.early_stopping(val_metrics['loss'])\n",
    "            if self.early_stopping.early_stop:\n",
    "                print(f\"Early stopping at epoch {epoch}\")\n",
    "                break\n",
    "        \n",
    "        self.model.load_state_dict(best_model)\n",
    "        return best_val_auc, best_attention_maps\n",
    "\n",
    "def compute_metrics(y_true, y_pred, y_prob):\n",
    "    \"\"\"Compute comprehensive set of metrics.\"\"\"\n",
    "    metrics = {\n",
    "        'accuracy': accuracy_score(y_true, y_pred),\n",
    "        'auc': roc_auc_score(y_true, y_prob[:, 1]) if y_prob.shape[1] == 2 else None,\n",
    "    }\n",
    "    \n",
    "    # Add precision-recall metrics\n",
    "    precision, recall, _ = precision_recall_curve(y_true, y_prob[:, 1])\n",
    "    metrics['avg_precision'] = average_precision_score(y_true, y_prob[:, 1])\n",
    "    \n",
    "    # Add confusion matrix based metrics\n",
    "    cm = confusion_matrix(y_true, y_pred)\n",
    "    metrics['specificity'] = cm[0,0] / (cm[0,0] + cm[0,1])\n",
    "    metrics['sensitivity'] = cm[1,1] / (cm[1,1] + cm[1,0])\n",
    "    \n",
    "    return metrics\n",
    "\n",
    "def visualize_training_history(history, save_path):\n",
    "    \"\"\"Plot training history with confidence intervals.\"\"\"\n",
    "    plt.style.use('seaborn-paper')\n",
    "    fig, axes = plt.subplots(2, 2, figsize=(15, 12))\n",
    "    \n",
    "    metrics = ['loss', 'accuracy', 'auc', 'avg_precision']\n",
    "    titles = ['Loss', 'Accuracy', 'ROC AUC', 'Average Precision']\n",
    "    \n",
    "    for ax, metric, title in zip(axes.flat, metrics, titles):\n",
    "        train_metric = history[f'train_{metric}']\n",
    "        val_metric = history[f'val_{metric}']\n",
    "        epochs = range(1, len(train_metric) + 1)\n",
    "        \n",
    "        ax.plot(epochs, train_metric, 'b-', label='Train')\n",
    "        ax.plot(epochs, val_metric, 'r-', label='Validation')\n",
    "        ax.set_title(title)\n",
    "        ax.set_xlabel('Epoch')\n",
    "        ax.set_ylabel(title)\n",
    "        ax.legend()\n",
    "        ax.grid(True, alpha=0.3)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, dpi=300, bbox_inches='tight')\n",
    "    plt.close()\n",
    "\n",
    "def visualize_attention_evolution(attention_maps, save_path):\n",
    "    \"\"\"Visualize how attention patterns evolve across layers.\"\"\"\n",
    "    plt.style.use('seaborn-paper')\n",
    "    num_layers = len(attention_maps)\n",
    "    \n",
    "    fig, axes = plt.subplots(1, num_layers, figsize=(5*num_layers, 4))\n",
    "    \n",
    "    for i, layer_maps in enumerate(attention_maps):\n",
    "        # Average attention weights across different scales\n",
    "        avg_attention = np.mean([m for m in layer_maps], axis=0)\n",
    "        \n",
    "        sns.heatmap(avg_attention, ax=axes[i], cmap='YlOrRd', \n",
    "                   xticklabels=20, yticklabels=20)\n",
    "        axes[i].set_title(f'Layer {i+1}')\n",
    "        \n",
    "    plt.suptitle('Evolution of Attention Patterns Across Layers', y=1.05)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, dpi=300, bbox_inches='tight')\n",
    "    plt.close()\n",
    "\n",
    "def visualize_multi_scale_attention(attention_maps, time_series, save_path):\n",
    "    \"\"\"Visualize attention at different scales.\"\"\"\n",
    "    plt.style.use('seaborn-paper')\n",
    "    num_scales = len(attention_maps[0])\n",
    "    \n",
    "    fig = plt.figure(figsize=(15, 3*num_scales))\n",
    "    gs = plt.GridSpec(num_scales, 1)\n",
    "    \n",
    "    for scale_idx in range(num_scales):\n",
    "        ax = fig.add_subplot(gs[scale_idx])\n",
    "        \n",
    "        # Plot time series\n",
    "        ax.plot(time_series, color='blue', alpha=0.6, label='Signal')\n",
    "        \n",
    "        # Plot attention weights\n",
    "        attention = attention_maps[0][scale_idx]  # Use first layer\n",
    "        attention_weights = attention.mean(axis=0)\n",
    "        \n",
    "        # Normalize attention weights for visualization\n",
    "        norm_weights = (attention_weights - attention_weights.min()) / \\\n",
    "                      (attention_weights.max() - attention_weights.min())\n",
    "                      \n",
    "        for i in range(len(time_series)):\n",
    "            ax.axvspan(i-0.5, i+0.5, color='red', alpha=0.1*norm_weights[i])\n",
    "            \n",
    "        ax.set_title(f'Scale {scale_idx + 1}')\n",
    "        ax.legend()\n",
    "        \n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, dpi=300, bbox_inches='tight')\n",
    "    plt.close()\n",
    "\n",
    "def plot_advanced_roc_curve(y_true, y_scores, fold_idx, save_path):\n",
    "    \"\"\"Plot ROC curve with confidence intervals.\"\"\"\n",
    "    plt.style.use('seaborn-paper')\n",
    "    fig, ax = plt.subplots(figsize=(8, 8))\n",
    "    \n",
    "    # Calculate ROC curve and AUC for each fold\n",
    "    fprs, tprs, roc_aucs = [], [], []\n",
    "    for i in range(len(fold_idx)):\n",
    "        fold_true = y_true[fold_idx[i]]\n",
    "        fold_score = y_scores[fold_idx[i]]\n",
    "        \n",
    "        fpr, tpr, _ = roc_curve(fold_true, fold_score[:, 1])\n",
    "        roc_auc = auc(fpr, tpr)\n",
    "        \n",
    "        fprs.append(fpr)\n",
    "        tprs.append(tpr)\n",
    "        roc_aucs.append(roc_auc)\n",
    "    \n",
    "    # Calculate mean ROC curve and AUC\n",
    "    mean_fpr = np.linspace(0, 1, 100)\n",
    "    mean_tpr = np.zeros_like(mean_fpr)\n",
    "    for fpr, tpr in zip(fprs, tprs):\n",
    "        mean_tpr += np.interp(mean_fpr, fpr, tpr)\n",
    "    mean_tpr /= len(fprs)\n",
    "    \n",
    "    mean_auc = np.mean(roc_aucs)\n",
    "    std_auc = np.std(roc_aucs)\n",
    "    \n",
    "    # Plot mean ROC curve\n",
    "    ax.plot(mean_fpr, mean_tpr, color='b',\n",
    "            label=f'Mean ROC (AUC = {mean_auc:.3f} ± {std_auc:.3f})',\n",
    "            lw=2, alpha=.8)\n",
    "    \n",
    "    # Plot confidence interval\n",
    "    tprs_upper = np.minimum(mean_tpr + np.std(tprs, axis=0), 1)\n",
    "    tprs_lower = np.maximum(mean_tpr - np.std(tprs, axis=0), 0)\n",
    "    ax.fill_between(mean_fpr, tprs_lower, tprs_upper, color='grey', alpha=.2,\n",
    "                    label=r'$\\pm$ 1 std. dev.')\n",
    "    \n",
    "    # Plot random chance line\n",
    "    ax.plot([0, 1], [0, 1], linestyle='--', lw=2, color='r',\n",
    "            label='Random Chance', alpha=.8)\n",
    "    \n",
    "    ax.set_xlim([-0.01, 1.01])\n",
    "    ax.set_ylim([-0.01, 1.01])\n",
    "    ax.set_xlabel('False Positive Rate', fontsize=12)\n",
    "    ax.set_ylabel('True Positive Rate', fontsize=12)\n",
    "    ax.set_title('Receiver Operating Characteristic (ROC) Curve', fontsize=14)\n",
    "    ax.legend(loc=\"lower right\", fontsize=10)\n",
    "    ax.grid(True, alpha=0.3)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, dpi=300, bbox_inches='tight')\n",
    "    plt.close()\n",
    "\n",
    "def generate_results_table(results_dict, dataset_name, save_path):\n",
    "    \"\"\"Generate comprehensive results table with confidence intervals.\"\"\"\n",
    "    metrics = ['accuracy', 'auc', 'specificity', 'sensitivity', 'avg_precision']\n",
    "    table_data = []\n",
    "    \n",
    "    for metric in metrics:\n",
    "        values = results_dict[metric]\n",
    "        mean = np.mean(values)\n",
    "        ci = stats.t.interval(0.95, len(values)-1, \n",
    "                            loc=mean, \n",
    "                            scale=stats.sem(values))\n",
    "        \n",
    "        table_data.append({\n",
    "            'Metric': metric.capitalize(),\n",
    "            'Mean': f\"{mean:.3f}\",\n",
    "            '95% CI': f\"({ci[0]:.3f}, {ci[1]:.3f})\"\n",
    "        })\n",
    "    \n",
    "    df = pd.DataFrame(table_data)\n",
    "    df.to_csv(save_path, index=False)\n",
    "    return df\n",
    "\n",
    "def run_experiment(dataset_name, config):\n",
    "    \"\"\"Run complete experiment pipeline with enhanced evaluation.\"\"\"\n",
    "    print(f\"\\n=== Dataset: {dataset_name} ===\")\n",
    "    \n",
    "    # Load and analyze dataset\n",
    "    X_all, y_all, unique_labels = load_dataset(dataset_name)\n",
    "    dataset_stats = analyze_dataset(X_all, y_all, dataset_name)\n",
    "    print(\"\\nDataset Statistics:\")\n",
    "    print(dataset_stats)\n",
    "    \n",
    "    # Convert to graphs with enhanced features\n",
    "    graphs = timeseries_to_graph(X_all)\n",
    "    labels = torch.tensor(y_all, dtype=torch.long)\n",
    "    \n",
    "    # Initialize results storage\n",
    "    results = defaultdict(list)\n",
    "    fold_predictions = []\n",
    "    fold_indices = []\n",
    "    attention_maps_collection = []\n",
    "    \n",
    "    # Cross-validation\n",
    "    skf = StratifiedKFold(n_splits=config.N_SPLITS, shuffle=True, random_state=42)\n",
    "    \n",
    "    for fold, (train_idx, test_idx) in enumerate(skf.split(graphs, labels), 1):\n",
    "        print(f\"\\nFold {fold}/{config.N_SPLITS}\")\n",
    "        \n",
    "        # Split data\n",
    "        train_graphs = [graphs[i] for i in train_idx]\n",
    "        test_graphs = [graphs[i] for i in test_idx]\n",
    "        train_labels = labels[train_idx]\n",
    "        test_labels = labels[test_idx]\n",
    "        \n",
    "        # Create dataloaders\n",
    "        train_loader = DataLoader(train_graphs, batch_size=config.BATCH_SIZE, shuffle=True)\n",
    "        test_loader = DataLoader(test_graphs, batch_size=config.BATCH_SIZE)\n",
    "        \n",
    "        # Initialize model and training components\n",
    "        model = EnhancedGNNClassifier(\n",
    "            input_dim=2, \n",
    "            hidden_dim=config.HIDDEN_DIM,\n",
    "            num_classes=len(unique_labels)\n",
    "        ).to(device)\n",
    "        \n",
    "        optimizer = optim.Adam(\n",
    "            model.parameters(), \n",
    "            lr=config.LR,\n",
    "            weight_decay=config.WEIGHT_DECAY\n",
    "        )\n",
    "        \n",
    "        criterion = nn.CrossEntropyLoss()\n",
    "        trainer = TrainingManager(model, optimizer, criterion, device, config)\n",
    "        \n",
    "        # Train model\n",
    "        best_val_auc, attention_maps = trainer.train(\n",
    "            train_loader, \n",
    "            test_loader,\n",
    "            config.MAX_EPOCHS\n",
    "        )\n",
    "        \n",
    "        # Store results\n",
    "        for k, v in trainer.history.items():\n",
    "            if k.startswith('val_'):\n",
    "                metric_name = k.split('_')[1]\n",
    "                results[metric_name].append(v[-1])\n",
    "        \n",
    "        attention_maps_collection.append(attention_maps)\n",
    "        fold_indices.append(test_idx)\n",
    "        \n",
    "        # Create visualizations for this fold\n",
    "        visualize_attention_evolution(\n",
    "            attention_maps,\n",
    "            os.path.join(config.PLOT_DIR, f'{dataset_name}_fold{fold}_attention_evolution.png')\n",
    "        )\n",
    "        \n",
    "        if fold == 1:  # Save training history for first fold\n",
    "            visualize_training_history(\n",
    "                trainer.history,\n",
    "                os.path.join(config.PLOT_DIR, f'{dataset_name}_training_history.png')\n",
    "            )\n",
    "    \n",
    "    # Generate final results\n",
    "    results_df = generate_results_table(\n",
    "        results,\n",
    "        dataset_name,\n",
    "        os.path.join(config.RESULTS_DIR, f'{dataset_name}_results.csv')\n",
    "    )\n",
    "    \n",
    "    print(\"\\nFinal Results:\")\n",
    "    print(results_df)\n",
    "    \n",
    "    return results, attention_maps_collection\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    print(f\"Using device: {device}\")\n",
    "    \n",
    "    for dataset_name in [\"ECG200\", \"ECGFiveDays\", \"TwoLeadECG\"]:\n",
    "        try:\n",
    "            results, attention_maps = run_experiment(dataset_name, config)\n",
    "        except Exception as e:\n",
    "            print(f\"Error processing dataset {dataset_name}: {str(e)}\")\n",
    "            continue\n",
    "    \n",
    "    print(\"\\nAll experiments completed. Results saved in:\", config.RESULTS_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f72aef36-9e72-4e0c-9523-a4fb73619693",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/salilpatel/.local/share/virtualenvs/salilpatel-ywhSAfpw/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cpu\n",
      "\n",
      "=== Dataset: ECG200 ===\n",
      "Number of classes: 2\n",
      "Total samples: 200 (Train: 100, Test: 100)\n",
      "Class -1: 67 samples\n",
      "Class 1: 133 samples\n",
      "\n",
      "Fold 1/5\n",
      "Epoch 0: Train Loss = 0.7026, Train Acc = 0.4609, Val Loss = 0.6961, Val Acc = 0.6562\n",
      "Epoch 10: Train Loss = 0.6858, Train Acc = 0.5703, Val Loss = 0.6897, Val Acc = 0.3438\n",
      "Epoch 20: Train Loss = 0.6707, Train Acc = 0.6875, Val Loss = 0.6589, Val Acc = 0.4062\n",
      "Epoch 30: Train Loss = 0.4710, Train Acc = 0.7969, Val Loss = 0.4841, Val Acc = 0.7500\n",
      "Epoch 40: Train Loss = 0.3532, Train Acc = 0.8438, Val Loss = 0.4912, Val Acc = 0.7500\n",
      "Epoch 50: Train Loss = 0.3623, Train Acc = 0.8516, Val Loss = 0.4403, Val Acc = 0.8125\n",
      "Early stopping at epoch 52\n",
      "Fold 1 - Test Accuracy: 0.9000\n",
      "\n",
      "Fold 2/5\n",
      "Epoch 0: Train Loss = 0.7153, Train Acc = 0.4766, Val Loss = 0.7008, Val Acc = 0.6562\n",
      "Epoch 10: Train Loss = 0.6914, Train Acc = 0.5078, Val Loss = 0.6899, Val Acc = 0.3438\n",
      "Epoch 20: Train Loss = 0.6463, Train Acc = 0.7188, Val Loss = 0.6127, Val Acc = 0.8125\n",
      "Epoch 30: Train Loss = 0.3546, Train Acc = 0.8516, Val Loss = 0.4329, Val Acc = 0.8750\n",
      "Epoch 40: Train Loss = 0.2934, Train Acc = 0.8984, Val Loss = 0.3659, Val Acc = 0.8438\n",
      "Epoch 50: Train Loss = 0.3045, Train Acc = 0.8828, Val Loss = 0.3403, Val Acc = 0.8750\n",
      "Epoch 60: Train Loss = 0.2347, Train Acc = 0.9219, Val Loss = 0.3460, Val Acc = 0.8750\n",
      "Epoch 70: Train Loss = 0.1983, Train Acc = 0.9141, Val Loss = 0.3024, Val Acc = 0.8750\n",
      "Early stopping at epoch 74\n",
      "Fold 2 - Test Accuracy: 0.6500\n",
      "\n",
      "Fold 3/5\n",
      "Epoch 0: Train Loss = 0.7135, Train Acc = 0.5547, Val Loss = 0.6934, Val Acc = 0.3438\n",
      "Epoch 10: Train Loss = 0.6755, Train Acc = 0.6484, Val Loss = 0.6807, Val Acc = 0.6875\n",
      "Epoch 20: Train Loss = 0.5089, Train Acc = 0.7891, Val Loss = 0.6281, Val Acc = 0.8125\n",
      "Epoch 30: Train Loss = 0.4532, Train Acc = 0.7969, Val Loss = 0.4581, Val Acc = 0.8125\n",
      "Epoch 40: Train Loss = 0.3935, Train Acc = 0.8516, Val Loss = 0.4048, Val Acc = 0.8125\n",
      "Epoch 50: Train Loss = 0.3159, Train Acc = 0.8672, Val Loss = 0.3561, Val Acc = 0.8125\n",
      "Epoch 60: Train Loss = 0.2926, Train Acc = 0.8672, Val Loss = 0.3472, Val Acc = 0.8750\n",
      "Epoch 70: Train Loss = 0.2647, Train Acc = 0.9062, Val Loss = 0.3057, Val Acc = 0.8750\n",
      "Epoch 80: Train Loss = 0.2058, Train Acc = 0.9219, Val Loss = 0.3120, Val Acc = 0.9062\n",
      "Early stopping at epoch 80\n",
      "Fold 3 - Test Accuracy: 0.7000\n",
      "\n",
      "Fold 4/5\n",
      "Epoch 0: Train Loss = 0.7062, Train Acc = 0.4609, Val Loss = 0.6931, Val Acc = 0.6562\n",
      "Epoch 10: Train Loss = 0.6918, Train Acc = 0.5078, Val Loss = 0.6910, Val Acc = 0.6562\n",
      "Epoch 20: Train Loss = 0.6823, Train Acc = 0.6328, Val Loss = 0.6704, Val Acc = 0.6562\n",
      "Epoch 30: Train Loss = 0.5867, Train Acc = 0.7344, Val Loss = 0.5327, Val Acc = 0.7500\n",
      "Epoch 40: Train Loss = 0.4926, Train Acc = 0.7266, Val Loss = 0.5381, Val Acc = 0.6875\n",
      "Epoch 50: Train Loss = 0.4578, Train Acc = 0.7500, Val Loss = 0.4108, Val Acc = 0.8125\n",
      "Epoch 60: Train Loss = 0.3842, Train Acc = 0.8281, Val Loss = 0.4114, Val Acc = 0.7812\n",
      "Epoch 70: Train Loss = 0.4386, Train Acc = 0.7578, Val Loss = 0.3855, Val Acc = 0.7812\n",
      "Epoch 80: Train Loss = 0.3422, Train Acc = 0.8594, Val Loss = 0.4629, Val Acc = 0.7188\n",
      "Early stopping at epoch 81\n",
      "Fold 4 - Test Accuracy: 0.8750\n",
      "\n",
      "Fold 5/5\n",
      "Epoch 0: Train Loss = 0.6881, Train Acc = 0.4922, Val Loss = 0.6916, Val Acc = 0.3438\n",
      "Epoch 10: Train Loss = 0.6948, Train Acc = 0.5234, Val Loss = 0.6897, Val Acc = 0.6562\n",
      "Epoch 20: Train Loss = 0.6696, Train Acc = 0.6953, Val Loss = 0.6466, Val Acc = 0.6562\n",
      "Epoch 30: Train Loss = 0.4943, Train Acc = 0.8047, Val Loss = 0.4509, Val Acc = 0.8438\n",
      "Epoch 40: Train Loss = 0.4528, Train Acc = 0.7891, Val Loss = 0.3858, Val Acc = 0.8438\n",
      "Epoch 50: Train Loss = 0.4166, Train Acc = 0.8359, Val Loss = 0.4024, Val Acc = 0.8125\n",
      "Early stopping at epoch 56\n",
      "Fold 5 - Test Accuracy: 0.9250\n",
      "\n",
      "Final Results:\n",
      "Validation Accuracy (mean ± std): 0.8750 ± 0.0280\n",
      "Test Accuracy: 0.8100 (95% CI: 0.7150-0.9100)\n",
      "Test AUC: 0.8775 (95% CI: 0.8285-0.9216)\n",
      "\n",
      "Classification Report:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "          -1       0.67      0.87      0.75        67\n",
      "           1       0.92      0.78      0.85       133\n",
      "\n",
      "    accuracy                           0.81       200\n",
      "   macro avg       0.79      0.82      0.80       200\n",
      "weighted avg       0.84      0.81      0.81       200\n",
      "\n",
      "\n",
      "=== Dataset: ECGFiveDays ===\n",
      "Number of classes: 2\n",
      "Total samples: 884 (Train: 23, Test: 861)\n",
      "Class 1: 442 samples\n",
      "Class 2: 442 samples\n",
      "\n",
      "Fold 1/5\n",
      "Epoch 0: Train Loss = 0.6899, Train Acc = 0.5424, Val Loss = 0.7135, Val Acc = 0.4681\n",
      "Epoch 10: Train Loss = 0.2375, Train Acc = 0.9028, Val Loss = 0.1378, Val Acc = 0.9362\n",
      "Epoch 20: Train Loss = 0.0773, Train Acc = 0.9841, Val Loss = 0.0622, Val Acc = 0.9858\n",
      "Epoch 30: Train Loss = 0.0426, Train Acc = 0.9912, Val Loss = 0.0549, Val Acc = 0.9858\n",
      "Epoch 40: Train Loss = 0.0278, Train Acc = 0.9912, Val Loss = 0.0466, Val Acc = 0.9787\n",
      "Epoch 50: Train Loss = 0.0134, Train Acc = 0.9982, Val Loss = 0.0185, Val Acc = 0.9929\n",
      "Early stopping at epoch 59\n",
      "Fold 1 - Test Accuracy: 0.9887\n",
      "\n",
      "Fold 2/5\n",
      "Epoch 0: Train Loss = 0.6979, Train Acc = 0.5000, Val Loss = 0.6872, Val Acc = 0.5390\n",
      "Epoch 10: Train Loss = 0.2995, Train Acc = 0.8693, Val Loss = 0.2812, Val Acc = 0.8936\n",
      "Epoch 20: Train Loss = 0.1212, Train Acc = 0.9629, Val Loss = 0.0871, Val Acc = 0.9787\n",
      "Epoch 30: Train Loss = 0.1192, Train Acc = 0.9558, Val Loss = 0.0581, Val Acc = 0.9858\n",
      "Epoch 40: Train Loss = 0.0358, Train Acc = 0.9912, Val Loss = 0.0306, Val Acc = 0.9929\n",
      "Epoch 50: Train Loss = 0.0269, Train Acc = 0.9947, Val Loss = 0.0293, Val Acc = 1.0000\n",
      "Epoch 60: Train Loss = 0.0140, Train Acc = 0.9982, Val Loss = 0.0187, Val Acc = 0.9929\n",
      "Epoch 70: Train Loss = 0.0285, Train Acc = 0.9912, Val Loss = 0.0480, Val Acc = 0.9787\n",
      "Early stopping at epoch 71\n",
      "Fold 2 - Test Accuracy: 0.9887\n",
      "\n",
      "Fold 3/5\n",
      "Epoch 0: Train Loss = 0.7074, Train Acc = 0.5106, Val Loss = 0.6851, Val Acc = 0.5532\n",
      "Epoch 10: Train Loss = 0.2357, Train Acc = 0.9081, Val Loss = 0.2100, Val Acc = 0.8936\n",
      "Epoch 20: Train Loss = 0.1140, Train Acc = 0.9647, Val Loss = 0.0825, Val Acc = 0.9858\n",
      "Epoch 30: Train Loss = 0.0498, Train Acc = 0.9912, Val Loss = 0.1304, Val Acc = 0.9433\n",
      "Epoch 40: Train Loss = 0.0409, Train Acc = 0.9841, Val Loss = 0.0231, Val Acc = 0.9929\n",
      "Epoch 50: Train Loss = 0.0430, Train Acc = 0.9806, Val Loss = 0.0248, Val Acc = 0.9858\n",
      "Epoch 60: Train Loss = 0.0169, Train Acc = 0.9929, Val Loss = 0.0291, Val Acc = 0.9858\n",
      "Early stopping at epoch 69\n",
      "Fold 3 - Test Accuracy: 1.0000\n",
      "\n",
      "Fold 4/5\n",
      "Epoch 0: Train Loss = 0.6957, Train Acc = 0.5106, Val Loss = 0.6891, Val Acc = 0.4823\n",
      "Epoch 10: Train Loss = 0.3614, Train Acc = 0.8428, Val Loss = 0.3053, Val Acc = 0.8865\n",
      "Epoch 20: Train Loss = 0.2967, Train Acc = 0.8799, Val Loss = 0.2210, Val Acc = 0.9220\n",
      "Epoch 30: Train Loss = 0.1837, Train Acc = 0.9329, Val Loss = 0.1191, Val Acc = 0.9716\n",
      "Epoch 40: Train Loss = 0.1049, Train Acc = 0.9647, Val Loss = 0.0488, Val Acc = 0.9929\n",
      "Epoch 50: Train Loss = 0.0401, Train Acc = 0.9912, Val Loss = 0.0111, Val Acc = 1.0000\n",
      "Epoch 60: Train Loss = 0.0200, Train Acc = 0.9947, Val Loss = 0.0100, Val Acc = 1.0000\n",
      "Epoch 70: Train Loss = 0.0681, Train Acc = 0.9770, Val Loss = 0.0341, Val Acc = 0.9929\n",
      "Early stopping at epoch 74\n",
      "Fold 4 - Test Accuracy: 0.9774\n",
      "\n",
      "Fold 5/5\n",
      "Epoch 0: Train Loss = 0.7007, Train Acc = 0.5009, Val Loss = 0.6898, Val Acc = 0.5248\n",
      "Epoch 10: Train Loss = 0.2396, Train Acc = 0.9101, Val Loss = 0.2087, Val Acc = 0.9291\n",
      "Epoch 20: Train Loss = 0.1312, Train Acc = 0.9577, Val Loss = 0.1147, Val Acc = 0.9716\n",
      "Epoch 30: Train Loss = 0.0671, Train Acc = 0.9806, Val Loss = 0.0472, Val Acc = 0.9858\n",
      "Epoch 40: Train Loss = 0.0377, Train Acc = 0.9894, Val Loss = 0.0342, Val Acc = 0.9929\n",
      "Epoch 50: Train Loss = 0.0455, Train Acc = 0.9894, Val Loss = 0.0243, Val Acc = 0.9929\n",
      "Epoch 60: Train Loss = 0.0763, Train Acc = 0.9718, Val Loss = 0.0575, Val Acc = 0.9858\n",
      "Early stopping at epoch 66\n",
      "Fold 5 - Test Accuracy: 1.0000\n",
      "\n",
      "Final Results:\n",
      "Validation Accuracy (mean ± std): 0.9972 ± 0.0035\n",
      "Test Accuracy: 0.9910 (95% CI: 0.9842-0.9977)\n",
      "Test AUC: 0.9987 (95% CI: 0.9965-0.9999)\n",
      "\n",
      "Classification Report:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           1       0.99      0.99      0.99       442\n",
      "           2       0.99      0.99      0.99       442\n",
      "\n",
      "    accuracy                           0.99       884\n",
      "   macro avg       0.99      0.99      0.99       884\n",
      "weighted avg       0.99      0.99      0.99       884\n",
      "\n",
      "\n",
      "=== Dataset: TwoLeadECG ===\n",
      "Number of classes: 2\n",
      "Total samples: 1162 (Train: 23, Test: 1139)\n",
      "Class 1: 581 samples\n",
      "Class 2: 581 samples\n",
      "\n",
      "Fold 1/5\n",
      "Epoch 0: Train Loss = 0.6933, Train Acc = 0.5376, Val Loss = 0.6864, Val Acc = 0.5297\n",
      "Epoch 10: Train Loss = 0.2858, Train Acc = 0.8925, Val Loss = 0.2422, Val Acc = 0.9081\n",
      "Epoch 20: Train Loss = 0.0747, Train Acc = 0.9731, Val Loss = 0.0605, Val Acc = 0.9730\n",
      "Epoch 30: Train Loss = 0.0332, Train Acc = 0.9919, Val Loss = 0.0384, Val Acc = 0.9946\n",
      "Epoch 40: Train Loss = 0.0375, Train Acc = 0.9866, Val Loss = 0.0283, Val Acc = 0.9946\n",
      "Early stopping at epoch 48\n",
      "Fold 1 - Test Accuracy: 0.9914\n",
      "\n",
      "Fold 2/5\n",
      "Epoch 0: Train Loss = 0.6966, Train Acc = 0.5255, Val Loss = 0.6925, Val Acc = 0.4811\n",
      "Epoch 10: Train Loss = 0.3106, Train Acc = 0.8616, Val Loss = 0.3663, Val Acc = 0.7838\n",
      "Epoch 20: Train Loss = 0.0844, Train Acc = 0.9731, Val Loss = 0.0841, Val Acc = 0.9730\n",
      "Epoch 30: Train Loss = 0.0555, Train Acc = 0.9852, Val Loss = 0.0343, Val Acc = 0.9892\n",
      "Epoch 40: Train Loss = 0.0592, Train Acc = 0.9839, Val Loss = 0.0285, Val Acc = 0.9892\n",
      "Epoch 50: Train Loss = 0.0501, Train Acc = 0.9852, Val Loss = 0.0235, Val Acc = 0.9838\n",
      "Epoch 60: Train Loss = 0.0281, Train Acc = 0.9919, Val Loss = 0.0184, Val Acc = 0.9946\n",
      "Epoch 70: Train Loss = 0.0450, Train Acc = 0.9892, Val Loss = 0.0094, Val Acc = 1.0000\n",
      "Epoch 80: Train Loss = 0.0319, Train Acc = 0.9892, Val Loss = 0.0139, Val Acc = 0.9946\n",
      "Early stopping at epoch 88\n",
      "Fold 2 - Test Accuracy: 0.9914\n",
      "\n",
      "Fold 3/5\n",
      "Epoch 0: Train Loss = 0.6958, Train Acc = 0.5228, Val Loss = 0.6989, Val Acc = 0.4570\n",
      "Epoch 10: Train Loss = 0.2358, Train Acc = 0.9234, Val Loss = 0.1770, Val Acc = 0.9355\n",
      "Epoch 20: Train Loss = 0.0710, Train Acc = 0.9758, Val Loss = 0.0338, Val Acc = 0.9946\n",
      "Epoch 30: Train Loss = 0.0556, Train Acc = 0.9825, Val Loss = 0.0198, Val Acc = 0.9946\n",
      "Epoch 40: Train Loss = 0.0165, Train Acc = 0.9946, Val Loss = 0.0099, Val Acc = 0.9946\n",
      "Epoch 50: Train Loss = 0.0358, Train Acc = 0.9892, Val Loss = 0.0756, Val Acc = 0.9839\n",
      "Early stopping at epoch 54\n",
      "Fold 3 - Test Accuracy: 0.9914\n",
      "\n",
      "Fold 4/5\n",
      "Epoch 0: Train Loss = 0.7022, Train Acc = 0.4987, Val Loss = 0.6958, Val Acc = 0.5161\n",
      "Epoch 10: Train Loss = 0.2535, Train Acc = 0.9019, Val Loss = 0.1894, Val Acc = 0.9516\n",
      "Epoch 20: Train Loss = 0.1045, Train Acc = 0.9677, Val Loss = 0.0505, Val Acc = 0.9892\n",
      "Epoch 30: Train Loss = 0.0417, Train Acc = 0.9906, Val Loss = 0.0590, Val Acc = 0.9839\n",
      "Epoch 40: Train Loss = 0.0607, Train Acc = 0.9812, Val Loss = 0.0544, Val Acc = 0.9677\n",
      "Early stopping at epoch 47\n",
      "Fold 4 - Test Accuracy: 0.9914\n",
      "\n",
      "Fold 5/5\n",
      "Epoch 0: Train Loss = 0.7031, Train Acc = 0.4825, Val Loss = 0.6937, Val Acc = 0.4892\n",
      "Epoch 10: Train Loss = 0.5299, Train Acc = 0.7675, Val Loss = 0.4482, Val Acc = 0.8548\n",
      "Epoch 20: Train Loss = 0.1576, Train Acc = 0.9422, Val Loss = 0.0801, Val Acc = 0.9892\n",
      "Epoch 30: Train Loss = 0.0493, Train Acc = 0.9879, Val Loss = 0.0581, Val Acc = 0.9892\n",
      "Early stopping at epoch 33\n",
      "Fold 5 - Test Accuracy: 0.9957\n",
      "\n",
      "Final Results:\n",
      "Validation Accuracy (mean ± std): 0.9968 ± 0.0026\n",
      "Test Accuracy: 0.9923 (95% CI: 0.9905-0.9931)\n",
      "Test AUC: 0.9990 (95% CI: 0.9976-0.9999)\n",
      "\n",
      "Classification Report:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           1       0.99      1.00      0.99       581\n",
      "           2       1.00      0.99      0.99       581\n",
      "\n",
      "    accuracy                           0.99      1162\n",
      "   macro avg       0.99      0.99      0.99      1162\n",
      "weighted avg       0.99      0.99      0.99      1162\n",
      "\n",
      "\n",
      "All experiments completed. Results saved in: ./plots\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import copy\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')  # For headless environments\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from sklearn.metrics import accuracy_score, confusion_matrix, classification_report, roc_curve, auc\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.loader import DataLoader\n",
    "from torch_geometric.nn import GATConv, global_mean_pool\n",
    "from torch_geometric.utils import add_self_loops\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import TensorDataset, DataLoader as TorchDataLoader\n",
    "from tslearn.datasets import UCR_UEA_datasets\n",
    "from scipy.stats import bootstrap\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# Configuration\n",
    "RANDOM_SEED = 42\n",
    "np.random.seed(RANDOM_SEED)\n",
    "torch.manual_seed(RANDOM_SEED)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed(RANDOM_SEED)\n",
    "\n",
    "MAX_EPOCHS = 100\n",
    "BATCH_SIZE = 32\n",
    "LR = 0.001\n",
    "N_SPLITS = 5\n",
    "HIDDEN_DIM = 64\n",
    "HEADS = 8\n",
    "DROPOUT = 0.2\n",
    "PLOT_DIR = \"./plots\"\n",
    "os.makedirs(PLOT_DIR, exist_ok=True)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "DATASETS = [\"ECG200\", \"ECGFiveDays\", \"TwoLeadECG\"]  # Example datasets\n",
    "\n",
    "class EarlyStopping:\n",
    "    def __init__(self, patience=10, min_delta=0):\n",
    "        self.patience = patience\n",
    "        self.min_delta = min_delta\n",
    "        self.counter = 0\n",
    "        self.best_loss = None\n",
    "        self.early_stop = False\n",
    "\n",
    "    def __call__(self, val_loss):\n",
    "        if self.best_loss is None:\n",
    "            self.best_loss = val_loss\n",
    "        elif val_loss > self.best_loss - self.min_delta:\n",
    "            self.counter += 1\n",
    "            if self.counter >= self.patience:\n",
    "                self.early_stop = True\n",
    "        else:\n",
    "            self.best_loss = val_loss\n",
    "            self.counter = 0\n",
    "\n",
    "def load_dataset(name):\n",
    "    \"\"\"Load UCR/UEA dataset, normalize, and return all data and labels.\"\"\"\n",
    "    ucr = UCR_UEA_datasets()\n",
    "    X_train, y_train, X_test, y_test = ucr.load_dataset(name)\n",
    "    \n",
    "    # Squeeze unnecessary dimensions for univariate TS\n",
    "    if X_train.ndim == 3 and X_train.shape[1] == 1:\n",
    "        X_train = X_train.squeeze(1)\n",
    "    if X_test.ndim == 3 and X_test.shape[1] == 1:\n",
    "        X_test = X_test.squeeze(1)\n",
    "    \n",
    "    unique_labels = np.unique(np.concatenate([y_train, y_test]))\n",
    "    label_to_idx = {lab: i for i, lab in enumerate(unique_labels)}\n",
    "    y_train = np.array([label_to_idx[lab] for lab in y_train], dtype=np.int64)\n",
    "    y_test = np.array([label_to_idx[lab] for lab in y_test], dtype=np.int64)\n",
    "    \n",
    "    # Normalize data\n",
    "    X_mean = X_train.mean()\n",
    "    X_std = X_train.std()\n",
    "    X_train = (X_train - X_mean) / (X_std + 1e-8)\n",
    "    X_test = (X_test - X_mean) / (X_std + 1e-8)\n",
    "    \n",
    "    X_all = np.concatenate([X_train, X_test], axis=0)\n",
    "    y_all = np.concatenate([y_train, y_test], axis=0)\n",
    "    \n",
    "    return X_all, y_all, unique_labels, (X_train, y_train, X_test, y_test)\n",
    "\n",
    "def timeseries_to_graph(X, window=5):\n",
    "    \"\"\"Convert time series to graph data by connecting each node to neighbors within a given window.\"\"\"\n",
    "    graphs = []\n",
    "    for i in range(X.shape[0]):\n",
    "        x_val = X[i]\n",
    "        length = x_val.shape[0]\n",
    "        positions = np.linspace(0, 1, length)\n",
    "        x_feat = np.column_stack([x_val, positions])\n",
    "        x_feat = torch.tensor(x_feat, dtype=torch.float32)\n",
    "        \n",
    "        edge_list = []\n",
    "        w = min(window, length)\n",
    "        for j in range(length):\n",
    "            for k in range(max(0, j-w), min(length, j+w+1)):\n",
    "                if j != k:\n",
    "                    edge_list.append([j, k])\n",
    "        \n",
    "        if edge_list:\n",
    "            edge_index = torch.tensor(edge_list, dtype=torch.long).t()\n",
    "        else:\n",
    "            edge_index = torch.empty((2, 0), dtype=torch.long)\n",
    "        \n",
    "        edge_index, _ = add_self_loops(edge_index, num_nodes=length)\n",
    "        data = Data(x=x_feat, edge_index=edge_index)\n",
    "        graphs.append(data)\n",
    "    return graphs\n",
    "\n",
    "class GATConvWithAlpha(GATConv):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.last_alpha = None\n",
    "        self._cached_edge_index = None\n",
    "\n",
    "    def forward(self, x, edge_index, return_attention_weights=False):\n",
    "        out, (edge_index_out, alpha) = super().forward(x, edge_index, return_attention_weights=True)\n",
    "        self.last_alpha = alpha\n",
    "        self._cached_edge_index = edge_index_out\n",
    "        return out\n",
    "\n",
    "class GNNTimeSeriesClassifier(nn.Module):\n",
    "    \"\"\"GAT-based classifier for time series represented as graphs.\"\"\"\n",
    "    def __init__(self, input_dim=2, hidden_dim=64, num_classes=2, heads=8, dropout=0.2):\n",
    "        super().__init__()\n",
    "        self.input_proj = nn.Linear(input_dim, hidden_dim)\n",
    "        self.conv1 = GATConvWithAlpha(hidden_dim, hidden_dim, heads=heads, dropout=dropout)\n",
    "        self.conv2 = GATConvWithAlpha(hidden_dim*heads, hidden_dim, heads=heads, dropout=dropout, concat=False)\n",
    "        \n",
    "        self.ln1 = nn.LayerNorm(hidden_dim * heads)\n",
    "        self.ln2 = nn.LayerNorm(hidden_dim)\n",
    "        \n",
    "        self.fc1 = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.fc2 = nn.Linear(hidden_dim, num_classes)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "        x = self.input_proj(x)\n",
    "        x = torch.relu(x)\n",
    "        \n",
    "        x1 = self.conv1(x, edge_index)\n",
    "        x1 = self.ln1(x1)\n",
    "        x1 = torch.relu(x1)\n",
    "        x1 = self.dropout(x1)\n",
    "        \n",
    "        x2 = self.conv2(x1, edge_index)\n",
    "        x2 = self.ln2(x2)\n",
    "        x2 = torch.relu(x2)\n",
    "        \n",
    "        x_pool = global_mean_pool(x2, batch)\n",
    "        x_pool = self.fc1(x_pool)\n",
    "        x_pool = torch.relu(x_pool)\n",
    "        x_pool = self.dropout(x_pool)\n",
    "        out = self.fc2(x_pool)\n",
    "        \n",
    "        return out\n",
    "\n",
    "    def get_attention_weights(self, data):\n",
    "        \"\"\"Obtain normalized node-level attention weights.\"\"\"\n",
    "        self.eval()\n",
    "        with torch.no_grad():\n",
    "            _ = self.forward(data)\n",
    "            alpha1 = self.conv1.last_alpha\n",
    "            edge_idx = self.conv1._cached_edge_index\n",
    "\n",
    "        num_nodes = data.x.size(0)\n",
    "        node_importance = torch.zeros(num_nodes, device=data.x.device)\n",
    "        \n",
    "        alpha_mean = alpha1.mean(dim=1)\n",
    "        for i in range(edge_idx.size(1)):\n",
    "            dst_node = edge_idx[1, i].item()\n",
    "            node_importance[dst_node] += alpha_mean[i].item()\n",
    "        \n",
    "        node_importance = node_importance / (node_importance.sum() + 1e-9)\n",
    "        return node_importance.cpu().numpy()\n",
    "\n",
    "def train_epoch(model, loader, criterion, optimizer, device):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    predictions = []\n",
    "    targets = []\n",
    "    \n",
    "    for batch in loader:\n",
    "        batch = batch.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        output = model(batch)\n",
    "        loss = criterion(output, batch.y.squeeze())\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        total_loss += loss.item()\n",
    "        predictions.extend(output.argmax(dim=1).cpu().numpy())\n",
    "        targets.extend(batch.y.cpu().numpy())\n",
    "    \n",
    "    return total_loss / len(loader), accuracy_score(targets, predictions)\n",
    "\n",
    "def evaluate(model, loader, criterion, device):\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    predictions = []\n",
    "    targets = []\n",
    "    outputs = []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for batch in loader:\n",
    "            batch = batch.to(device)\n",
    "            output = model(batch)\n",
    "            loss = criterion(output, batch.y.squeeze())\n",
    "            \n",
    "            total_loss += loss.item()\n",
    "            predictions.extend(output.argmax(dim=1).cpu().numpy())\n",
    "            targets.extend(batch.y.cpu().numpy())\n",
    "            outputs.extend(torch.softmax(output, dim=1).cpu().numpy())\n",
    "    \n",
    "    return (\n",
    "        total_loss / len(loader),\n",
    "        accuracy_score(targets, predictions),\n",
    "        np.array(predictions),\n",
    "        np.array(targets),\n",
    "        np.array(outputs)\n",
    "    )\n",
    "\n",
    "def train_model(model, train_loader, val_loader, criterion, optimizer, device, max_epochs=100):\n",
    "    early_stopping = EarlyStopping(patience=10)\n",
    "    best_val_acc = 0\n",
    "    best_model = None\n",
    "    \n",
    "    for epoch in range(max_epochs):\n",
    "        train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)\n",
    "        val_loss, val_acc, _, _, _ = evaluate(model, val_loader, criterion, device)\n",
    "        \n",
    "        if val_acc > best_val_acc:\n",
    "            best_val_acc = val_acc\n",
    "            best_model = copy.deepcopy(model.state_dict())\n",
    "        \n",
    "        # Optional: print every 10 epochs\n",
    "        if epoch % 10 == 0:\n",
    "            print(f'Epoch {epoch}: Train Loss = {train_loss:.4f}, Train Acc = {train_acc:.4f}, '\n",
    "                  f'Val Loss = {val_loss:.4f}, Val Acc = {val_acc:.4f}')\n",
    "        \n",
    "        early_stopping(val_loss)\n",
    "        if early_stopping.early_stop:\n",
    "            print(f\"Early stopping at epoch {epoch}\")\n",
    "            break\n",
    "    \n",
    "    model.load_state_dict(best_model)\n",
    "    return model, best_val_acc\n",
    "\n",
    "def bootstrap_confidence_interval(data, stat_func=np.mean, confidence_level=0.95, n_resamples=10000):\n",
    "    \"\"\"Compute bootstrap confidence interval for a given statistic.\"\"\"\n",
    "    res = bootstrap((data,), stat_func, confidence_level=confidence_level, n_resamples=n_resamples, method='basic')\n",
    "    return stat_func(data), res.confidence_interval\n",
    "\n",
    "def plot_confusion_matrix(cm, classes, title, save_path):\n",
    "    plt.figure(figsize=(8, 6))\n",
    "    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=classes, yticklabels=classes)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.ylabel('True Label', fontsize=14)\n",
    "    plt.xlabel('Predicted Label', fontsize=14)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='pdf')\n",
    "    plt.close()\n",
    "\n",
    "def plot_roc_with_ci(fprs, tprs, title, save_path):\n",
    "    \"\"\"Plot mean ROC curve and 95% CI band from multiple folds.\"\"\"\n",
    "    # Interpolate TPRs at common FPR points\n",
    "    mean_fpr = np.linspace(0, 1, 100)\n",
    "    interp_tprs = []\n",
    "    for fpr, tpr in zip(fprs, tprs):\n",
    "        interp_tpr = np.interp(mean_fpr, fpr, tpr)\n",
    "        interp_tprs.append(interp_tpr)\n",
    "    interp_tprs = np.array(interp_tprs)\n",
    "    \n",
    "    mean_tpr = interp_tprs.mean(axis=0)\n",
    "    std_tpr = interp_tprs.std(axis=0)\n",
    "    tpr_upper = np.minimum(mean_tpr + 1.96 * std_tpr, 1)\n",
    "    tpr_lower = np.maximum(mean_tpr - 1.96 * std_tpr, 0)\n",
    "    mean_auc = auc(mean_fpr, mean_tpr)\n",
    "    \n",
    "    plt.figure(figsize=(8,6))\n",
    "    plt.plot(mean_fpr, mean_tpr, color='darkorange', lw=2, label=f'Mean ROC (AUC = {mean_auc:.2f})')\n",
    "    plt.fill_between(mean_fpr, tpr_lower, tpr_upper, color='grey', alpha=0.2, label='95% CI')\n",
    "    plt.plot([0,1],[0,1], color='navy', lw=2, linestyle='--')\n",
    "    plt.xlim([0,1])\n",
    "    plt.ylim([0,1.05])\n",
    "    plt.xlabel('False Positive Rate', fontsize=14)\n",
    "    plt.ylabel('True Positive Rate', fontsize=14)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.legend(loc='lower right')\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='pdf')\n",
    "    plt.close()\n",
    "\n",
    "def visualize_attention(series, attention_weights, title, save_path):\n",
    "    \"\"\"Visualize time series with attention weights as a color overlay.\"\"\"\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    plt.plot(series, label='Time Series', color='blue')\n",
    "    \n",
    "    # Normalize attention weights for visualization\n",
    "    att_norm = (attention_weights - attention_weights.min()) / (attention_weights.max() - attention_weights.min() + 1e-9)\n",
    "\n",
    "    # Create a colormap (e.g., Reds)\n",
    "    cmap = plt.cm.Reds\n",
    "    for i in range(len(series)):\n",
    "        color = cmap(att_norm[i])\n",
    "        plt.axvspan(i-0.5, i+0.5, color=color, alpha=0.3)\n",
    "    \n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.xlabel('Time Step', fontsize=14)\n",
    "    plt.ylabel('Normalized Value', fontsize=14)\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='pdf')\n",
    "    plt.close()\n",
    "\n",
    "def plot_example_timeseries(X, y, unique_labels, dataset_name, save_path):\n",
    "    \"\"\"Plot a few example time series from each class for qualitative inspection.\"\"\"\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    num_classes = len(unique_labels)\n",
    "    examples_per_class = min(3, len(X)//num_classes)\n",
    "    \n",
    "    for i, cls in enumerate(unique_labels):\n",
    "        class_indices = np.where(y == i)[0]\n",
    "        chosen = np.random.choice(class_indices, size=examples_per_class, replace=False)\n",
    "        for j, idx in enumerate(chosen):\n",
    "            plt.subplot(num_classes, examples_per_class, i*examples_per_class + j + 1)\n",
    "            plt.plot(X[idx], color='blue')\n",
    "            plt.title(f'Class: {cls}', fontsize=12)\n",
    "            plt.xticks([])\n",
    "            plt.yticks([])\n",
    "    plt.suptitle(f'{dataset_name} - Example Time Series', fontsize=16)\n",
    "    plt.tight_layout(rect=[0,0,1,0.95])\n",
    "    plt.savefig(save_path, format='pdf')\n",
    "    plt.close()\n",
    "\n",
    "def run_experiment(dataset_name):\n",
    "    print(f\"\\n=== Dataset: {dataset_name} ===\")\n",
    "    \n",
    "    # Load and preprocess data\n",
    "    X_all, y_all, unique_labels, (X_train_orig, y_train_orig, X_test_orig, y_test_orig) = load_dataset(dataset_name)\n",
    "    graphs = timeseries_to_graph(X_all)\n",
    "    labels = torch.tensor(y_all, dtype=torch.long)\n",
    "    num_classes = len(unique_labels)\n",
    "\n",
    "    # Print dataset details\n",
    "    print(f\"Number of classes: {num_classes}\")\n",
    "    print(f\"Total samples: {len(X_all)} (Train: {len(X_train_orig)}, Test: {len(X_test_orig)})\")\n",
    "    class_counts = np.bincount(y_all)\n",
    "    for i, lab in enumerate(unique_labels):\n",
    "        print(f\"Class {lab}: {class_counts[i]} samples\")\n",
    "\n",
    "    # Plot example time series\n",
    "    plot_example_timeseries(X_all, y_all, unique_labels, dataset_name, \n",
    "                            save_path=os.path.join(PLOT_DIR, f\"{dataset_name}_example_ts.pdf\"))\n",
    "\n",
    "    all_val_accs = []\n",
    "    all_test_accs = []\n",
    "    all_test_preds = []\n",
    "    all_test_targets = []\n",
    "    all_test_probs = []\n",
    "    fprs_list = []\n",
    "    tprs_list = []\n",
    "\n",
    "    skf = StratifiedKFold(n_splits=N_SPLITS, shuffle=True, random_state=RANDOM_SEED)\n",
    "    for fold, (train_idx, test_idx) in enumerate(skf.split(graphs, labels), 1):\n",
    "        print(f\"\\nFold {fold}/{N_SPLITS}\")\n",
    "        \n",
    "        val_size = int(0.2 * len(train_idx))\n",
    "        train_indices = train_idx[:-val_size]\n",
    "        val_indices = train_idx[-val_size:]\n",
    "        \n",
    "        # Create datasets\n",
    "        train_dataset = []\n",
    "        for i in train_indices:\n",
    "            graph = graphs[i]\n",
    "            graph.y = labels[i]\n",
    "            train_dataset.append(graph)\n",
    "            \n",
    "        val_dataset = []\n",
    "        for i in val_indices:\n",
    "            graph = graphs[i]\n",
    "            graph.y = labels[i]\n",
    "            val_dataset.append(graph)\n",
    "            \n",
    "        test_dataset = []\n",
    "        for i in test_idx:\n",
    "            graph = graphs[i]\n",
    "            graph.y = labels[i]\n",
    "            test_dataset.append(graph)\n",
    "        \n",
    "        train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "        val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
    "        test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
    "        \n",
    "        model = GNNTimeSeriesClassifier(\n",
    "            input_dim=2,\n",
    "            hidden_dim=HIDDEN_DIM,\n",
    "            num_classes=num_classes,\n",
    "            heads=HEADS,\n",
    "            dropout=DROPOUT\n",
    "        ).to(device)\n",
    "        \n",
    "        optimizer = optim.Adam(model.parameters(), lr=LR, weight_decay=1e-4)\n",
    "        \n",
    "        # Class weighting\n",
    "        train_labels_np = labels[train_indices].numpy()\n",
    "        class_counts_train = np.bincount(train_labels_np)\n",
    "        class_weights = torch.FloatTensor(1.0 / class_counts_train).to(device)\n",
    "        criterion = nn.CrossEntropyLoss(weight=class_weights)\n",
    "        \n",
    "        model, val_acc = train_model(\n",
    "            model=model,\n",
    "            train_loader=train_loader,\n",
    "            val_loader=val_loader,\n",
    "            criterion=criterion,\n",
    "            optimizer=optimizer,\n",
    "            device=device,\n",
    "            max_epochs=MAX_EPOCHS\n",
    "        )\n",
    "        \n",
    "        test_loss, test_acc, predictions, targets, probabilities = evaluate(\n",
    "            model=model,\n",
    "            loader=test_loader,\n",
    "            criterion=criterion,\n",
    "            device=device\n",
    "        )\n",
    "        \n",
    "        # Compute AUC if binary classification\n",
    "        if num_classes == 2:\n",
    "            fpr, tpr, _ = roc_curve(targets, probabilities[:, 1])\n",
    "            fprs_list.append(fpr)\n",
    "            tprs_list.append(tpr)\n",
    "        \n",
    "        all_val_accs.append(val_acc)\n",
    "        all_test_accs.append(test_acc)\n",
    "        all_test_preds.extend(predictions)\n",
    "        all_test_targets.extend(targets)\n",
    "        all_test_probs.extend(probabilities)\n",
    "        \n",
    "        print(f\"Fold {fold} - Test Accuracy: {test_acc:.4f}\")\n",
    "        \n",
    "        # Visualize attention weights for first test sample (if available)\n",
    "        if len(test_dataset) > 0:\n",
    "            example_data = test_dataset[0].to(device)\n",
    "            attention_weights = model.get_attention_weights(example_data)\n",
    "            example_series = example_data.x[:, 0].cpu().numpy()\n",
    "            \n",
    "            visualize_attention(\n",
    "                series=example_series,\n",
    "                attention_weights=attention_weights,\n",
    "                title=f\"{dataset_name} - Fold {fold} Attention Visualization\",\n",
    "                save_path=os.path.join(PLOT_DIR, f\"{dataset_name}_fold{fold}_attention.pdf\")\n",
    "            )\n",
    "    \n",
    "    all_test_preds = np.array(all_test_preds)\n",
    "    all_test_targets = np.array(all_test_targets)\n",
    "    all_test_probs = np.array(all_test_probs)\n",
    "\n",
    "    # Compute bootstrap CIs for accuracy and AUC\n",
    "    mean_acc, ci_acc = bootstrap_confidence_interval(all_test_accs)\n",
    "    print(\"\\nFinal Results:\")\n",
    "    print(f\"Validation Accuracy (mean ± std): {np.mean(all_val_accs):.4f} ± {np.std(all_val_accs):.4f}\")\n",
    "    print(f\"Test Accuracy: {mean_acc:.4f} (95% CI: {ci_acc.low:.4f}-{ci_acc.high:.4f})\")\n",
    "\n",
    "    if num_classes == 2:\n",
    "        # Compute AUC across all folds (using combined predictions)\n",
    "        combined_fpr, combined_tpr, _ = roc_curve(all_test_targets, all_test_probs[:,1])\n",
    "        combined_auc = auc(combined_fpr, combined_tpr)\n",
    "        \n",
    "        # Bootstrap AUC CI\n",
    "        def auc_stat(data):\n",
    "            # data: (targets, probs)\n",
    "            targets_ = data[0].astype(int)\n",
    "            probs_ = data[1]\n",
    "            fpr_, tpr_, _ = roc_curve(targets_, probs_)\n",
    "            return auc(fpr_, tpr_)\n",
    "        \n",
    "        auc_samples = []\n",
    "        rng = np.random.default_rng(RANDOM_SEED)\n",
    "        for _ in range(10000):\n",
    "            idx = rng.integers(0, len(all_test_targets), len(all_test_targets))\n",
    "            auc_samples.append(auc_stat((all_test_targets[idx], all_test_probs[idx,1])))\n",
    "        auc_samples = np.array(auc_samples)\n",
    "        mean_auc_ = np.mean(auc_samples)\n",
    "        low_auc_ = np.percentile(auc_samples, 2.5)\n",
    "        high_auc_ = np.percentile(auc_samples, 97.5)\n",
    "        \n",
    "        print(f\"Test AUC: {mean_auc_:.4f} (95% CI: {low_auc_:.4f}-{high_auc_:.4f})\")\n",
    "    \n",
    "    # Plot confusion matrix\n",
    "    cm = confusion_matrix(all_test_targets, all_test_preds)\n",
    "    plot_confusion_matrix(\n",
    "        cm=cm,\n",
    "        classes=[str(l) for l in unique_labels],\n",
    "        title=f\"{dataset_name} Confusion Matrix\",\n",
    "        save_path=os.path.join(PLOT_DIR, f\"{dataset_name}_confusion_matrix.pdf\")\n",
    "    )\n",
    "    \n",
    "    # Plot aggregated ROC curve with CIs (for binary classification)\n",
    "    if num_classes == 2 and len(fprs_list) == N_SPLITS:\n",
    "        plot_roc_with_ci(\n",
    "            fprs=fprs_list,\n",
    "            tprs=tprs_list,\n",
    "            title=f\"{dataset_name} ROC Curve (with 95% CI)\",\n",
    "            save_path=os.path.join(PLOT_DIR, f\"{dataset_name}_roc_curve.pdf\")\n",
    "        )\n",
    "    \n",
    "    # Print classification report\n",
    "    print(\"\\nClassification Report:\")\n",
    "    print(classification_report(all_test_targets, all_test_preds, target_names=[str(l) for l in unique_labels]))\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    print(f\"Using device: {device}\")\n",
    "    for dataset_name in DATASETS:\n",
    "        try:\n",
    "            run_experiment(dataset_name)\n",
    "        except Exception as e:\n",
    "            print(f\"Error processing dataset {dataset_name}: {str(e)}\")\n",
    "            continue\n",
    "    \n",
    "    print(\"\\nAll experiments completed. Results saved in:\", PLOT_DIR)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d1d985d6-4e5d-4fa7-86eb-6832ba765b00",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cpu\n",
      "\n",
      "=== Dataset: ECG200 ===\n",
      "Number of classes: 2\n",
      "Total samples: 200 (Train: 100, Test: 100)\n",
      "Class -1: 67 samples\n",
      "Class 1: 133 samples\n",
      "\n",
      "Fold 1/5\n",
      "Epoch 0: Train Loss = 0.7026, Train Acc = 0.4609, Val Loss = 0.6961, Val Acc = 0.6562\n",
      "Epoch 10: Train Loss = 0.6858, Train Acc = 0.5703, Val Loss = 0.6897, Val Acc = 0.3438\n",
      "Epoch 20: Train Loss = 0.6707, Train Acc = 0.6875, Val Loss = 0.6589, Val Acc = 0.4062\n",
      "Epoch 30: Train Loss = 0.4710, Train Acc = 0.7969, Val Loss = 0.4841, Val Acc = 0.7500\n",
      "Epoch 40: Train Loss = 0.3532, Train Acc = 0.8438, Val Loss = 0.4912, Val Acc = 0.7500\n",
      "Epoch 50: Train Loss = 0.3623, Train Acc = 0.8516, Val Loss = 0.4403, Val Acc = 0.8125\n",
      "Early stopping at epoch 52\n",
      "Fold 1 - Test Accuracy: 0.9000, Balanced Acc: 0.9066, MCC: 0.7917\n",
      "Error processing dataset ECG200: invalid index of a 0-dim tensor. Use `tensor.item()` in Python or `tensor.item<T>()` in C++ to convert a 0-dim tensor to a number\n",
      "\n",
      "=== Dataset: ECGFiveDays ===\n",
      "Number of classes: 2\n",
      "Total samples: 884 (Train: 23, Test: 861)\n",
      "Class 1: 442 samples\n",
      "Class 2: 442 samples\n",
      "\n",
      "Fold 1/5\n",
      "Epoch 0: Train Loss = 0.7063, Train Acc = 0.4823, Val Loss = 0.6886, Val Acc = 0.5319\n",
      "Epoch 10: Train Loss = 0.2651, Train Acc = 0.8905, Val Loss = 0.2184, Val Acc = 0.9007\n",
      "Epoch 20: Train Loss = 0.1691, Train Acc = 0.9276, Val Loss = 0.0976, Val Acc = 0.9787\n",
      "Epoch 30: Train Loss = 0.0906, Train Acc = 0.9682, Val Loss = 0.0572, Val Acc = 0.9716\n",
      "Epoch 40: Train Loss = 0.0293, Train Acc = 0.9947, Val Loss = 0.0332, Val Acc = 0.9929\n",
      "Epoch 50: Train Loss = 0.1485, Train Acc = 0.9505, Val Loss = 0.0863, Val Acc = 0.9645\n",
      "Early stopping at epoch 52\n",
      "Fold 1 - Test Accuracy: 0.9944, Balanced Acc: 0.9943, MCC: 0.9888\n",
      "Error processing dataset ECGFiveDays: invalid index of a 0-dim tensor. Use `tensor.item()` in Python or `tensor.item<T>()` in C++ to convert a 0-dim tensor to a number\n",
      "\n",
      "=== Dataset: TwoLeadECG ===\n",
      "Number of classes: 2\n",
      "Total samples: 1162 (Train: 23, Test: 1139)\n",
      "Class 1: 581 samples\n",
      "Class 2: 581 samples\n",
      "\n",
      "Fold 1/5\n",
      "Epoch 0: Train Loss = 0.7008, Train Acc = 0.4933, Val Loss = 0.7136, Val Acc = 0.4703\n",
      "Epoch 10: Train Loss = 0.3375, Train Acc = 0.8522, Val Loss = 0.1982, Val Acc = 0.9622\n",
      "Epoch 20: Train Loss = 0.0711, Train Acc = 0.9785, Val Loss = 0.0298, Val Acc = 1.0000\n",
      "Epoch 30: Train Loss = 0.0258, Train Acc = 0.9933, Val Loss = 0.0132, Val Acc = 0.9946\n",
      "Epoch 40: Train Loss = 0.1083, Train Acc = 0.9597, Val Loss = 0.0242, Val Acc = 0.9946\n",
      "Early stopping at epoch 47\n",
      "Fold 1 - Test Accuracy: 0.9957, Balanced Acc: 0.9957, MCC: 0.9915\n",
      "Error processing dataset TwoLeadECG: invalid index of a 0-dim tensor. Use `tensor.item()` in Python or `tensor.item<T>()` in C++ to convert a 0-dim tensor to a number\n",
      "\n",
      "All experiments completed. Results saved in: ./plots\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import copy\n",
    "import warnings\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')  # For headless environments\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from scipy.stats import bootstrap\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from sklearn.metrics import (accuracy_score, confusion_matrix, classification_report, roc_curve, auc,\n",
    "                             balanced_accuracy_score, matthews_corrcoef)\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.loader import DataLoader\n",
    "from torch_geometric.nn import GATConv, global_mean_pool\n",
    "from torch_geometric.utils import add_self_loops\n",
    "from tslearn.datasets import UCR_UEA_datasets\n",
    "from captum.attr import Saliency\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# Configuration\n",
    "RANDOM_SEED = 42\n",
    "np.random.seed(RANDOM_SEED)\n",
    "torch.manual_seed(RANDOM_SEED)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed(RANDOM_SEED)\n",
    "\n",
    "MAX_EPOCHS = 100\n",
    "BATCH_SIZE = 32\n",
    "LR = 0.001\n",
    "N_SPLITS = 5\n",
    "HIDDEN_DIM = 64\n",
    "HEADS = 8\n",
    "DROPOUT = 0.2\n",
    "PLOT_DIR = \"./plots\"\n",
    "os.makedirs(PLOT_DIR, exist_ok=True)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "DATASETS = [\"ECG200\", \"ECGFiveDays\", \"TwoLeadECG\"]  # Example datasets\n",
    "\n",
    "class EarlyStopping:\n",
    "    def __init__(self, patience=10, min_delta=0):\n",
    "        self.patience = patience\n",
    "        self.min_delta = min_delta\n",
    "        self.counter = 0\n",
    "        self.best_loss = None\n",
    "        self.early_stop = False\n",
    "\n",
    "    def __call__(self, val_loss):\n",
    "        if self.best_loss is None:\n",
    "            self.best_loss = val_loss\n",
    "        elif val_loss > self.best_loss - self.min_delta:\n",
    "            self.counter += 1\n",
    "            if self.counter >= self.patience:\n",
    "                self.early_stop = True\n",
    "        else:\n",
    "            self.best_loss = val_loss\n",
    "            self.counter = 0\n",
    "\n",
    "def load_dataset(name):\n",
    "    \"\"\"Load UCR/UEA dataset, normalize, and return all data and labels.\"\"\"\n",
    "    ucr = UCR_UEA_datasets()\n",
    "    X_train, y_train, X_test, y_test = ucr.load_dataset(name)\n",
    "    \n",
    "    # Squeeze unnecessary dimensions for univariate TS\n",
    "    if X_train.ndim == 3 and X_train.shape[1] == 1:\n",
    "        X_train = X_train.squeeze(1)\n",
    "    if X_test.ndim == 3 and X_test.shape[1] == 1:\n",
    "        X_test = X_test.squeeze(1)\n",
    "    \n",
    "    unique_labels = np.unique(np.concatenate([y_train, y_test]))\n",
    "    label_to_idx = {lab: i for i, lab in enumerate(unique_labels)}\n",
    "    y_train = np.array([label_to_idx[lab] for lab in y_train], dtype=np.int64)\n",
    "    y_test = np.array([label_to_idx[lab] for lab in y_test], dtype=np.int64)\n",
    "    \n",
    "    # Normalize data\n",
    "    X_mean = X_train.mean()\n",
    "    X_std = X_train.std()\n",
    "    X_train = (X_train - X_mean) / (X_std + 1e-8)\n",
    "    X_test = (X_test - X_mean) / (X_std + 1e-8)\n",
    "    \n",
    "    X_all = np.concatenate([X_train, X_test], axis=0)\n",
    "    y_all = np.concatenate([y_train, y_test], axis=0)\n",
    "    \n",
    "    return X_all, y_all, unique_labels, (X_train, y_train, X_test, y_test)\n",
    "\n",
    "def timeseries_to_graph(X, window=5):\n",
    "    \"\"\"Convert time series to graph data by connecting each node to neighbors within a given window.\"\"\"\n",
    "    graphs = []\n",
    "    for i in range(X.shape[0]):\n",
    "        x_val = X[i]\n",
    "        length = x_val.shape[0]\n",
    "        positions = np.linspace(0, 1, length)\n",
    "        x_feat = np.column_stack([x_val, positions])\n",
    "        x_feat = torch.tensor(x_feat, dtype=torch.float32)\n",
    "        \n",
    "        edge_list = []\n",
    "        w = min(window, length)\n",
    "        for j in range(length):\n",
    "            for k in range(max(0, j-w), min(length, j+w+1)):\n",
    "                if j != k:\n",
    "                    edge_list.append([j, k])\n",
    "        \n",
    "        if edge_list:\n",
    "            edge_index = torch.tensor(edge_list, dtype=torch.long).t()\n",
    "        else:\n",
    "            edge_index = torch.empty((2, 0), dtype=torch.long)\n",
    "        \n",
    "        edge_index, _ = add_self_loops(edge_index, num_nodes=length)\n",
    "        data = Data(x=x_feat, edge_index=edge_index)\n",
    "        graphs.append(data)\n",
    "    return graphs\n",
    "\n",
    "class GATConvWithAlpha(GATConv):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.last_alpha = None\n",
    "        self._cached_edge_index = None\n",
    "\n",
    "    def forward(self, x, edge_index, return_attention_weights=False):\n",
    "        out, (edge_index_out, alpha) = super().forward(x, edge_index, return_attention_weights=True)\n",
    "        self.last_alpha = alpha\n",
    "        self._cached_edge_index = edge_index_out\n",
    "        return out\n",
    "\n",
    "class GNNTimeSeriesClassifier(nn.Module):\n",
    "    \"\"\"GAT-based classifier for time series represented as graphs.\"\"\"\n",
    "    def __init__(self, input_dim=2, hidden_dim=64, num_classes=2, heads=8, dropout=0.2):\n",
    "        super().__init__()\n",
    "        self.input_proj = nn.Linear(input_dim, hidden_dim)\n",
    "        self.conv1 = GATConvWithAlpha(hidden_dim, hidden_dim, heads=heads, dropout=dropout)\n",
    "        self.conv2 = GATConvWithAlpha(hidden_dim*heads, hidden_dim, heads=heads, dropout=dropout, concat=False)\n",
    "        \n",
    "        self.ln1 = nn.LayerNorm(hidden_dim * heads)\n",
    "        self.ln2 = nn.LayerNorm(hidden_dim)\n",
    "        \n",
    "        self.fc1 = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.fc2 = nn.Linear(hidden_dim, num_classes)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "        x = self.input_proj(x)\n",
    "        x = torch.relu(x)\n",
    "        \n",
    "        x1 = self.conv1(x, edge_index)\n",
    "        x1 = self.ln1(x1)\n",
    "        x1 = torch.relu(x1)\n",
    "        x1 = self.dropout(x1)\n",
    "        \n",
    "        x2 = self.conv2(x1, edge_index)\n",
    "        x2 = self.ln2(x2)\n",
    "        x2 = torch.relu(x2)\n",
    "        \n",
    "        x_pool = global_mean_pool(x2, batch)\n",
    "        x_pool = self.fc1(x_pool)\n",
    "        x_pool = torch.relu(x_pool)\n",
    "        x_pool = self.dropout(x_pool)\n",
    "        out = self.fc2(x_pool)\n",
    "        \n",
    "        return out\n",
    "\n",
    "    def get_attention_weights(self, data):\n",
    "        \"\"\"Obtain normalized node-level attention weights from the first GAT layer.\"\"\"\n",
    "        self.eval()\n",
    "        with torch.no_grad():\n",
    "            _ = self.forward(data)\n",
    "            alpha1 = self.conv1.last_alpha\n",
    "            edge_idx = self.conv1._cached_edge_index\n",
    "\n",
    "        num_nodes = data.x.size(0)\n",
    "        node_importance = torch.zeros(num_nodes, device=data.x.device)\n",
    "        \n",
    "        alpha_mean = alpha1.mean(dim=1)\n",
    "        for i in range(edge_idx.size(1)):\n",
    "            dst_node = edge_idx[1, i].item()\n",
    "            node_importance[dst_node] += alpha_mean[i].item()\n",
    "        \n",
    "        node_importance = node_importance / (node_importance.sum() + 1e-9)\n",
    "        return node_importance.cpu().numpy()\n",
    "\n",
    "def train_epoch(model, loader, criterion, optimizer, device):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    predictions = []\n",
    "    targets = []\n",
    "    \n",
    "    for batch in loader:\n",
    "        batch = batch.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        output = model(batch)\n",
    "        loss = criterion(output, batch.y.squeeze())\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        total_loss += loss.item()\n",
    "        predictions.extend(output.argmax(dim=1).cpu().numpy())\n",
    "        targets.extend(batch.y.cpu().numpy())\n",
    "    \n",
    "    return total_loss / len(loader), accuracy_score(targets, predictions)\n",
    "\n",
    "def evaluate(model, loader, criterion, device):\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    predictions = []\n",
    "    targets = []\n",
    "    outputs = []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for batch in loader:\n",
    "            batch = batch.to(device)\n",
    "            output = model(batch)\n",
    "            loss = criterion(output, batch.y.squeeze())\n",
    "            \n",
    "            total_loss += loss.item()\n",
    "            predictions.extend(output.argmax(dim=1).cpu().numpy())\n",
    "            targets.extend(batch.y.cpu().numpy())\n",
    "            outputs.extend(torch.softmax(output, dim=1).cpu().numpy())\n",
    "    \n",
    "    return (\n",
    "        total_loss / len(loader),\n",
    "        accuracy_score(targets, predictions),\n",
    "        np.array(predictions),\n",
    "        np.array(targets),\n",
    "        np.array(outputs)\n",
    "    )\n",
    "\n",
    "def train_model(model, train_loader, val_loader, criterion, optimizer, device, max_epochs=100):\n",
    "    early_stopping = EarlyStopping(patience=10)\n",
    "    best_val_acc = 0\n",
    "    best_model = None\n",
    "    \n",
    "    for epoch in range(max_epochs):\n",
    "        train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)\n",
    "        val_loss, val_acc, _, _, _ = evaluate(model, val_loader, criterion, device)\n",
    "        \n",
    "        if val_acc > best_val_acc:\n",
    "            best_val_acc = val_acc\n",
    "            best_model = copy.deepcopy(model.state_dict())\n",
    "        \n",
    "        if epoch % 10 == 0:\n",
    "            print(f'Epoch {epoch}: Train Loss = {train_loss:.4f}, Train Acc = {train_acc:.4f}, '\n",
    "                  f'Val Loss = {val_loss:.4f}, Val Acc = {val_acc:.4f}')\n",
    "        \n",
    "        early_stopping(val_loss)\n",
    "        if early_stopping.early_stop:\n",
    "            print(f\"Early stopping at epoch {epoch}\")\n",
    "            break\n",
    "    \n",
    "    model.load_state_dict(best_model)\n",
    "    return model, best_val_acc\n",
    "\n",
    "def bootstrap_confidence_interval(data, stat_func=np.mean, confidence_level=0.95, n_resamples=10000):\n",
    "    \"\"\"Compute bootstrap confidence interval for a given statistic.\"\"\"\n",
    "    res = bootstrap((data,), stat_func, confidence_level=confidence_level, n_resamples=n_resamples, method='basic')\n",
    "    return stat_func(data), res.confidence_interval\n",
    "\n",
    "def plot_confusion_matrix(cm, classes, title, save_path):\n",
    "    plt.figure(figsize=(8, 6))\n",
    "    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=classes, yticklabels=classes)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.ylabel('True Label', fontsize=14)\n",
    "    plt.xlabel('Predicted Label', fontsize=14)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='pdf')\n",
    "    plt.close()\n",
    "\n",
    "def plot_roc_with_ci(fprs, tprs, title, save_path):\n",
    "    \"\"\"Plot mean ROC curve and 95% CI band from multiple folds.\"\"\"\n",
    "    mean_fpr = np.linspace(0, 1, 100)\n",
    "    interp_tprs = []\n",
    "    for fpr, tpr in zip(fprs, tprs):\n",
    "        interp_tpr = np.interp(mean_fpr, fpr, tpr)\n",
    "        interp_tprs.append(interp_tpr)\n",
    "    interp_tprs = np.array(interp_tprs)\n",
    "    \n",
    "    mean_tpr = interp_tprs.mean(axis=0)\n",
    "    std_tpr = interp_tprs.std(axis=0)\n",
    "    tpr_upper = np.minimum(mean_tpr + 1.96 * std_tpr, 1)\n",
    "    tpr_lower = np.maximum(mean_tpr - 1.96 * std_tpr, 0)\n",
    "    mean_auc = auc(mean_fpr, mean_tpr)\n",
    "    \n",
    "    plt.figure(figsize=(8,6))\n",
    "    plt.plot(mean_fpr, mean_tpr, color='darkorange', lw=2, label=f'Mean ROC (AUC = {mean_auc:.2f})')\n",
    "    plt.fill_between(mean_fpr, tpr_lower, tpr_upper, color='grey', alpha=0.2, label='95% CI')\n",
    "    plt.plot([0,1],[0,1], color='navy', lw=2, linestyle='--')\n",
    "    plt.xlim([0,1])\n",
    "    plt.ylim([0,1.05])\n",
    "    plt.xlabel('False Positive Rate', fontsize=14)\n",
    "    plt.ylabel('True Positive Rate', fontsize=14)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.legend(loc='lower right')\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='pdf')\n",
    "    plt.close()\n",
    "\n",
    "def visualize_attention(series, attention_weights, title, save_path):\n",
    "    \"\"\"Visualize time series with attention weights as a color overlay.\"\"\"\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    plt.plot(series, label='Time Series', color='blue')\n",
    "    att_norm = (attention_weights - attention_weights.min()) / (attention_weights.max() - attention_weights.min() + 1e-9)\n",
    "    cmap = plt.cm.Reds\n",
    "    for i in range(len(series)):\n",
    "        color = cmap(att_norm[i])\n",
    "        plt.axvspan(i-0.5, i+0.5, color=color, alpha=0.3)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.xlabel('Time Step', fontsize=14)\n",
    "    plt.ylabel('Normalized Value', fontsize=14)\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='pdf')\n",
    "    plt.close()\n",
    "\n",
    "def plot_example_timeseries(X, y, unique_labels, dataset_name, save_path):\n",
    "    \"\"\"Plot a few example time series from each class for qualitative inspection.\"\"\"\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    num_classes = len(unique_labels)\n",
    "    examples_per_class = min(3, len(X)//num_classes)\n",
    "    \n",
    "    for i, cls in enumerate(unique_labels):\n",
    "        class_indices = np.where(y == i)[0]\n",
    "        chosen = np.random.choice(class_indices, size=examples_per_class, replace=False)\n",
    "        for j, idx in enumerate(chosen):\n",
    "            plt.subplot(num_classes, examples_per_class, i*examples_per_class + j + 1)\n",
    "            plt.plot(X[idx], color='blue')\n",
    "            plt.title(f'Class: {cls}', fontsize=12)\n",
    "            plt.xticks([])\n",
    "            plt.yticks([])\n",
    "    plt.suptitle(f'{dataset_name} - Example Time Series', fontsize=16)\n",
    "    plt.tight_layout(rect=[0,0,1,0.95])\n",
    "    plt.savefig(save_path, format='pdf')\n",
    "    plt.close()\n",
    "\n",
    "def plot_graph_with_attention(data, attention_weights, title, save_path):\n",
    "    \"\"\"Visualize graph structure with nodes colored by attention weights.\n",
    "    Here we position nodes along x-axis (time) and y-axis as their original value.\"\"\"\n",
    "    x_coords = data.x[:,1].cpu().numpy() * len(data.x)  # position scaled to length\n",
    "    y_vals = data.x[:,0].cpu().numpy()\n",
    "    att_norm = (attention_weights - attention_weights.min()) / (attention_weights.max() - attention_weights.min() + 1e-9)\n",
    "\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    # Draw edges\n",
    "    edge_index = data.edge_index.cpu().numpy()\n",
    "    for i in range(edge_index.shape[1]):\n",
    "        src, dst = edge_index[:, i]\n",
    "        plt.plot([x_coords[src], x_coords[dst]], [y_vals[src], y_vals[dst]], color='lightgray', linewidth=1, alpha=0.5)\n",
    "\n",
    "    # Draw nodes\n",
    "    plt.scatter(x_coords, y_vals, c=att_norm, cmap='Reds', s=50, edgecolors='black')\n",
    "    plt.colorbar(label='Attention Weight (normalized)')\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.xlabel('Time Step (scaled)')\n",
    "    plt.ylabel('Normalized Value')\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='pdf')\n",
    "    plt.close()\n",
    "\n",
    "def plot_attention_distribution(att_weights_all, title, save_path):\n",
    "    \"\"\"Plot the distribution (histogram) of attention weights aggregated across samples.\"\"\"\n",
    "    plt.figure(figsize=(8,6))\n",
    "    plt.hist(att_weights_all, bins=30, color='steelblue', edgecolor='black', alpha=0.7)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.xlabel('Attention Weight')\n",
    "    plt.ylabel('Frequency')\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='pdf')\n",
    "    plt.close()\n",
    "\n",
    "def plot_average_attention_profile(all_att_arrays, title, save_path):\n",
    "    \"\"\"Plot mean and confidence interval of attention weights across samples at each timestep.\"\"\"\n",
    "    # all_att_arrays: list of arrays of shape (num_nodes,)\n",
    "    # We need to align them by length. Let's assume all have same length for simplicity.\n",
    "    lengths = [len(a) for a in all_att_arrays]\n",
    "    min_len = min(lengths)  # if not equal, truncate\n",
    "    trimmed = np.array([a[:min_len] for a in all_att_arrays])\n",
    "    mean_att = trimmed.mean(axis=0)\n",
    "    std_att = trimmed.std(axis=0)\n",
    "\n",
    "    timesteps = np.arange(min_len)\n",
    "    plt.figure(figsize=(10,6))\n",
    "    plt.plot(timesteps, mean_att, color='red', linewidth=2, label='Mean Attention')\n",
    "    plt.fill_between(timesteps, mean_att - 1.96*std_att, mean_att + 1.96*std_att, color='pink', alpha=0.3, label='95% CI')\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.xlabel('Time Step')\n",
    "    plt.ylabel('Attention Weight')\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='pdf')\n",
    "    plt.close()\n",
    "\n",
    "def compute_saliency(model, data, device):\n",
    "    \"\"\"Compute saliency (gradient-based) using Captum to understand feature importance.\"\"\"\n",
    "    model.eval()\n",
    "    data = data.to(device)\n",
    "    # We consider input node features as baseline\n",
    "    saliency = Saliency(model_forward_wrapper(model))\n",
    "    baseline = data.x.clone().detach().requires_grad_(True)\n",
    "    attributions = saliency.attribute(data.x.unsqueeze(0), target=int(data.y[0].item()), additional_forward_args=(data.edge_index, data.batch))\n",
    "    # attributions shape: [1, num_nodes, num_features]\n",
    "    return attributions.squeeze(0).cpu().numpy()\n",
    "\n",
    "def model_forward_wrapper(model):\n",
    "    \"\"\"Wrapper for model forward to use Captum with node-level graph data.\"\"\"\n",
    "    def forward(x, edge_index, batch):\n",
    "        # Rebuild a Data object\n",
    "        data = Data(x=x[0], edge_index=edge_index, batch=batch)\n",
    "        out = model(data)\n",
    "        return out\n",
    "    return forward\n",
    "\n",
    "def run_experiment(dataset_name):\n",
    "    print(f\"\\n=== Dataset: {dataset_name} ===\")\n",
    "    \n",
    "    # Load and preprocess data\n",
    "    X_all, y_all, unique_labels, (X_train_orig, y_train_orig, X_test_orig, y_test_orig) = load_dataset(dataset_name)\n",
    "    graphs = timeseries_to_graph(X_all)\n",
    "    labels = torch.tensor(y_all, dtype=torch.long)\n",
    "    num_classes = len(unique_labels)\n",
    "\n",
    "    # Print dataset details\n",
    "    print(f\"Number of classes: {num_classes}\")\n",
    "    print(f\"Total samples: {len(X_all)} (Train: {len(X_train_orig)}, Test: {len(X_test_orig)})\")\n",
    "    class_counts = np.bincount(y_all)\n",
    "    for i, lab in enumerate(unique_labels):\n",
    "        print(f\"Class {lab}: {class_counts[i]} samples\")\n",
    "\n",
    "    # Plot example time series\n",
    "    plot_example_timeseries(X_all, y_all, unique_labels, dataset_name, \n",
    "                            save_path=os.path.join(PLOT_DIR, f\"{dataset_name}_example_ts.pdf\"))\n",
    "\n",
    "    all_val_accs = []\n",
    "    all_test_accs = []\n",
    "    all_bal_accs = []\n",
    "    all_mccs = []\n",
    "    all_test_preds = []\n",
    "    all_test_targets = []\n",
    "    all_test_probs = []\n",
    "    fprs_list = []\n",
    "    tprs_list = []\n",
    "    all_test_attentions = []\n",
    "\n",
    "    skf = StratifiedKFold(n_splits=N_SPLITS, shuffle=True, random_state=RANDOM_SEED)\n",
    "    for fold, (train_idx, test_idx) in enumerate(skf.split(graphs, labels), 1):\n",
    "        print(f\"\\nFold {fold}/{N_SPLITS}\")\n",
    "        \n",
    "        val_size = int(0.2 * len(train_idx))\n",
    "        train_indices = train_idx[:-val_size]\n",
    "        val_indices = train_idx[-val_size:]\n",
    "        \n",
    "        # Create datasets\n",
    "        train_dataset = []\n",
    "        for i in train_indices:\n",
    "            graph = graphs[i]\n",
    "            graph.y = labels[i]\n",
    "            train_dataset.append(graph)\n",
    "            \n",
    "        val_dataset = []\n",
    "        for i in val_indices:\n",
    "            graph = graphs[i]\n",
    "            graph.y = labels[i]\n",
    "            val_dataset.append(graph)\n",
    "            \n",
    "        test_dataset = []\n",
    "        for i in test_idx:\n",
    "            graph = graphs[i]\n",
    "            graph.y = labels[i]\n",
    "            test_dataset.append(graph)\n",
    "        \n",
    "        train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "        val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
    "        test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
    "        \n",
    "        model = GNNTimeSeriesClassifier(\n",
    "            input_dim=2,\n",
    "            hidden_dim=HIDDEN_DIM,\n",
    "            num_classes=num_classes,\n",
    "            heads=HEADS,\n",
    "            dropout=DROPOUT\n",
    "        ).to(device)\n",
    "        \n",
    "        optimizer = optim.Adam(model.parameters(), lr=LR, weight_decay=1e-4)\n",
    "        \n",
    "        # Class weighting\n",
    "        train_labels_np = labels[train_indices].numpy()\n",
    "        class_counts_train = np.bincount(train_labels_np)\n",
    "        class_weights = torch.FloatTensor(1.0 / class_counts_train).to(device)\n",
    "        criterion = nn.CrossEntropyLoss(weight=class_weights)\n",
    "        \n",
    "        model, val_acc = train_model(\n",
    "            model=model,\n",
    "            train_loader=train_loader,\n",
    "            val_loader=val_loader,\n",
    "            criterion=criterion,\n",
    "            optimizer=optimizer,\n",
    "            device=device,\n",
    "            max_epochs=MAX_EPOCHS\n",
    "        )\n",
    "        \n",
    "        test_loss, test_acc, predictions, targets, probabilities = evaluate(\n",
    "            model=model,\n",
    "            loader=test_loader,\n",
    "            criterion=criterion,\n",
    "            device=device\n",
    "        )\n",
    "\n",
    "        bal_acc = balanced_accuracy_score(targets, predictions)\n",
    "        mcc = matthews_corrcoef(targets, predictions)\n",
    "\n",
    "        # Compute AUC if binary classification\n",
    "        if num_classes == 2:\n",
    "            fpr, tpr, _ = roc_curve(targets, probabilities[:, 1])\n",
    "            fprs_list.append(fpr)\n",
    "            tprs_list.append(tpr)\n",
    "        \n",
    "        # Store metrics\n",
    "        all_val_accs.append(val_acc)\n",
    "        all_test_accs.append(test_acc)\n",
    "        all_bal_accs.append(bal_acc)\n",
    "        all_mccs.append(mcc)\n",
    "        all_test_preds.extend(predictions)\n",
    "        all_test_targets.extend(targets)\n",
    "        all_test_probs.extend(probabilities)\n",
    "\n",
    "        print(f\"Fold {fold} - Test Accuracy: {test_acc:.4f}, Balanced Acc: {bal_acc:.4f}, MCC: {mcc:.4f}\")\n",
    "\n",
    "        # Attention visualization on a test sample\n",
    "        if len(test_dataset) > 0:\n",
    "            example_data = test_dataset[0].clone().to(device)\n",
    "            attention_weights = model.get_attention_weights(example_data)\n",
    "            example_series = example_data.x[:, 0].cpu().numpy()\n",
    "            visualize_attention(\n",
    "                series=example_series,\n",
    "                attention_weights=attention_weights,\n",
    "                title=f\"{dataset_name} - Fold {fold} Attention on TS\",\n",
    "                save_path=os.path.join(PLOT_DIR, f\"{dataset_name}_fold{fold}_attention_ts.pdf\")\n",
    "            )\n",
    "            plot_graph_with_attention(\n",
    "                data=example_data.cpu(),\n",
    "                attention_weights=attention_weights,\n",
    "                title=f\"{dataset_name} - Fold {fold} Graph with Node Attention\",\n",
    "                save_path=os.path.join(PLOT_DIR, f\"{dataset_name}_fold{fold}_attention_graph.pdf\")\n",
    "            )\n",
    "            \n",
    "            # Also compute saliency for interpretation\n",
    "            saliency_vals = compute_saliency(model, example_data.clone(), device)\n",
    "            # Plot saliency as a heatmap\n",
    "            plt.figure(figsize=(10,4))\n",
    "            plt.imshow(saliency_vals.T, aspect='auto', cmap='coolwarm')\n",
    "            plt.colorbar(label='Saliency')\n",
    "            plt.title(f\"{dataset_name} - Fold {fold} Saliency Heatmap\")\n",
    "            plt.xlabel('Node Index')\n",
    "            plt.ylabel('Feature (0:value, 1:position)')\n",
    "            plt.tight_layout()\n",
    "            plt.savefig(os.path.join(PLOT_DIR, f\"{dataset_name}_fold{fold}_saliency.pdf\"), format='pdf')\n",
    "            plt.close()\n",
    "\n",
    "        # Collect attention weights from all test samples for aggregated analysis\n",
    "        test_att_weights = []\n",
    "        for d in test_dataset:\n",
    "            d = d.to(device)\n",
    "            aw = model.get_attention_weights(d)\n",
    "            test_att_weights.append(aw)\n",
    "        all_test_attentions.extend(test_att_weights)\n",
    "\n",
    "    all_test_preds = np.array(all_test_preds)\n",
    "    all_test_targets = np.array(all_test_targets)\n",
    "    all_test_probs = np.array(all_test_probs)\n",
    "\n",
    "    # Compute bootstrap CIs for accuracy and AUC\n",
    "    mean_acc, ci_acc = bootstrap_confidence_interval(all_test_accs)\n",
    "    mean_bal_acc, ci_bal_acc = bootstrap_confidence_interval(all_bal_accs)\n",
    "    mean_mcc, ci_mcc = bootstrap_confidence_interval(all_mccs)\n",
    "\n",
    "    print(\"\\nFinal Results:\")\n",
    "    print(f\"Validation Accuracy (mean ± std): {np.mean(all_val_accs):.4f} ± {np.std(all_val_accs):.4f}\")\n",
    "    print(f\"Test Accuracy: {mean_acc:.4f} (95% CI: {ci_acc.low:.4f}-{ci_acc.high:.4f})\")\n",
    "    print(f\"Balanced Accuracy: {mean_bal_acc:.4f} (95% CI: {ci_bal_acc.low:.4f}-{ci_bal_acc.high:.4f})\")\n",
    "    print(f\"MCC: {mean_mcc:.4f} (95% CI: {ci_mcc.low:.4f}-{ci_mcc.high:.4f})\")\n",
    "\n",
    "    if num_classes == 2:\n",
    "        # Compute AUC across all samples\n",
    "        combined_fpr, combined_tpr, _ = roc_curve(all_test_targets, all_test_probs[:,1])\n",
    "        combined_auc = auc(combined_fpr, combined_tpr)\n",
    "        \n",
    "        # Bootstrap AUC\n",
    "        def auc_stat(data):\n",
    "            targets_, probs_ = data\n",
    "            fpr_, tpr_, _ = roc_curve(targets_.astype(int), probs_)\n",
    "            return auc(fpr_, tpr_)\n",
    "\n",
    "        auc_samples = []\n",
    "        rng = np.random.default_rng(RANDOM_SEED)\n",
    "        for _ in range(10000):\n",
    "            idx = rng.integers(0, len(all_test_targets), len(all_test_targets))\n",
    "            auc_samples.append(auc_stat((all_test_targets[idx], all_test_probs[idx,1])))\n",
    "        auc_samples = np.array(auc_samples)\n",
    "        mean_auc_ = np.mean(auc_samples)\n",
    "        low_auc_ = np.percentile(auc_samples, 2.5)\n",
    "        high_auc_ = np.percentile(auc_samples, 97.5)\n",
    "        \n",
    "        print(f\"Test AUC: {mean_auc_:.4f} (95% CI: {low_auc_:.4f}-{high_auc_:.4f})\")\n",
    "\n",
    "    # Plot confusion matrix\n",
    "    cm = confusion_matrix(all_test_targets, all_test_preds)\n",
    "    plot_confusion_matrix(\n",
    "        cm=cm,\n",
    "        classes=[str(l) for l in unique_labels],\n",
    "        title=f\"{dataset_name} Confusion Matrix\",\n",
    "        save_path=os.path.join(PLOT_DIR, f\"{dataset_name}_confusion_matrix.pdf\")\n",
    "    )\n",
    "    \n",
    "    # Plot aggregated ROC curve with CIs (for binary classification)\n",
    "    if num_classes == 2 and len(fprs_list) == N_SPLITS:\n",
    "        plot_roc_with_ci(\n",
    "            fprs=fprs_list,\n",
    "            tprs=tprs_list,\n",
    "            title=f\"{dataset_name} ROC Curve (with 95% CI)\",\n",
    "            save_path=os.path.join(PLOT_DIR, f\"{dataset_name}_roc_curve.pdf\")\n",
    "        )\n",
    "\n",
    "    # Classification report\n",
    "    print(\"\\nClassification Report:\")\n",
    "    print(classification_report(all_test_targets, all_test_preds, target_names=[str(l) for l in unique_labels]))\n",
    "\n",
    "    # Attention distribution and average profile\n",
    "    all_test_attentions_arr = np.concatenate(all_test_attentions)\n",
    "    plot_attention_distribution(all_test_attentions_arr,\n",
    "                                title=f\"{dataset_name} - Attention Weight Distribution\",\n",
    "                                save_path=os.path.join(PLOT_DIR, f\"{dataset_name}_attention_distribution.pdf\"))\n",
    "\n",
    "    plot_average_attention_profile(all_test_attentions,\n",
    "                                   title=f\"{dataset_name} - Average Attention Profile\",\n",
    "                                   save_path=os.path.join(PLOT_DIR, f\"{dataset_name}_average_attention_profile.pdf\"))\n",
    "\n",
    "    # Save metrics to CSV\n",
    "    results_dict = {\n",
    "        'Val_Accuracy': all_val_accs,\n",
    "        'Test_Accuracy': all_test_accs,\n",
    "        'Balanced_Accuracy': all_bal_accs,\n",
    "        'MCC': all_mccs\n",
    "    }\n",
    "    if num_classes == 2:\n",
    "        results_dict['Test_AUC'] = [auc_stat((all_test_targets, all_test_probs[:,1]))]*len(all_test_accs)\n",
    "\n",
    "    # Mean and CI also stored\n",
    "    summary_dict = {\n",
    "        'Mean_Accuracy': [mean_acc],\n",
    "        'CI_Accuracy_Low': [ci_acc.low],\n",
    "        'CI_Accuracy_High': [ci_acc.high],\n",
    "        'Mean_Balanced_Accuracy': [mean_bal_acc],\n",
    "        'CI_Bal_Acc_Low': [ci_bal_acc.low],\n",
    "        'CI_Bal_Acc_High': [ci_bal_acc.high],\n",
    "        'Mean_MCC': [mean_mcc],\n",
    "        'CI_MCC_Low': [ci_mcc.low],\n",
    "        'CI_MCC_High': [ci_mcc.high]\n",
    "    }\n",
    "\n",
    "    if num_classes == 2:\n",
    "        summary_dict['Mean_AUC'] = [mean_auc_]\n",
    "        summary_dict['CI_AUC_Low'] = [low_auc_]\n",
    "        summary_dict['CI_AUC_High'] = [high_auc_]\n",
    "\n",
    "    import pandas as pd\n",
    "    pd.DataFrame(results_dict).to_csv(os.path.join(PLOT_DIR, f\"{dataset_name}_fold_metrics.csv\"), index=False)\n",
    "    pd.DataFrame(summary_dict).to_csv(os.path.join(PLOT_DIR, f\"{dataset_name}_summary_metrics.csv\"), index=False)\n",
    "    \n",
    "    print(f\"Results and plots saved to {PLOT_DIR}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    print(f\"Using device: {device}\")\n",
    "    for dataset_name in DATASETS:\n",
    "        try:\n",
    "            run_experiment(dataset_name)\n",
    "        except Exception as e:\n",
    "            print(f\"Error processing dataset {dataset_name}: {str(e)}\")\n",
    "            continue\n",
    "    \n",
    "    print(\"\\nAll experiments completed. Results saved in:\", PLOT_DIR)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d93bcce9-769e-416d-a6a6-90bc859b9165",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cpu\n",
      "\n",
      "=== Dataset: ECG200 ===\n",
      "Number of classes: 2\n",
      "Total samples: 200 (Train: 100, Test: 100)\n",
      "Class -1: 67 samples\n",
      "Class 1: 133 samples\n",
      "\n",
      "Fold 1/5\n",
      "Epoch 0: Train Loss = 0.7026, Train Acc = 0.4609, Val Loss = 0.6961, Val Acc = 0.6562\n",
      "Epoch 10: Train Loss = 0.6858, Train Acc = 0.5703, Val Loss = 0.6897, Val Acc = 0.3438\n",
      "Epoch 20: Train Loss = 0.6707, Train Acc = 0.6875, Val Loss = 0.6590, Val Acc = 0.4062\n",
      "Epoch 30: Train Loss = 0.4736, Train Acc = 0.7969, Val Loss = 0.4828, Val Acc = 0.7500\n",
      "Epoch 40: Train Loss = 0.3536, Train Acc = 0.8594, Val Loss = 0.4859, Val Acc = 0.7500\n",
      "Epoch 50: Train Loss = 0.3882, Train Acc = 0.8281, Val Loss = 0.4613, Val Acc = 0.8125\n",
      "Early stopping at epoch 52\n",
      "Fold 1 - Test Accuracy: 0.9000, Balanced Acc: 0.9066, MCC: 0.7917\n",
      "\n",
      "Fold 2/5\n",
      "Epoch 0: Train Loss = 0.7153, Train Acc = 0.4766, Val Loss = 0.7008, Val Acc = 0.6562\n",
      "Epoch 10: Train Loss = 0.6914, Train Acc = 0.5078, Val Loss = 0.6899, Val Acc = 0.3438\n",
      "Epoch 20: Train Loss = 0.6464, Train Acc = 0.7031, Val Loss = 0.6127, Val Acc = 0.8125\n",
      "Epoch 30: Train Loss = 0.3515, Train Acc = 0.8516, Val Loss = 0.4255, Val Acc = 0.8750\n",
      "Epoch 40: Train Loss = 0.2953, Train Acc = 0.8906, Val Loss = 0.3753, Val Acc = 0.8438\n",
      "Epoch 50: Train Loss = 0.3097, Train Acc = 0.8828, Val Loss = 0.3553, Val Acc = 0.8125\n",
      "Epoch 60: Train Loss = 0.2373, Train Acc = 0.9062, Val Loss = 0.3726, Val Acc = 0.8750\n",
      "Epoch 70: Train Loss = 0.1957, Train Acc = 0.9141, Val Loss = 0.3085, Val Acc = 0.9062\n",
      "Early stopping at epoch 74\n",
      "Fold 2 - Test Accuracy: 0.6500, Balanced Acc: 0.6978, MCC: 0.3852\n",
      "\n",
      "Fold 3/5\n",
      "Epoch 0: Train Loss = 0.7135, Train Acc = 0.5547, Val Loss = 0.6934, Val Acc = 0.3438\n",
      "Epoch 10: Train Loss = 0.6755, Train Acc = 0.6484, Val Loss = 0.6807, Val Acc = 0.6875\n",
      "Epoch 20: Train Loss = 0.5103, Train Acc = 0.7891, Val Loss = 0.6295, Val Acc = 0.8125\n",
      "Epoch 30: Train Loss = 0.4534, Train Acc = 0.7969, Val Loss = 0.4438, Val Acc = 0.8125\n",
      "Epoch 40: Train Loss = 0.3817, Train Acc = 0.8594, Val Loss = 0.3860, Val Acc = 0.7812\n",
      "Epoch 50: Train Loss = 0.3341, Train Acc = 0.8438, Val Loss = 0.3487, Val Acc = 0.8125\n",
      "Epoch 60: Train Loss = 0.2791, Train Acc = 0.8672, Val Loss = 0.3491, Val Acc = 0.8750\n",
      "Epoch 70: Train Loss = 0.2718, Train Acc = 0.9062, Val Loss = 0.2887, Val Acc = 0.9062\n",
      "Epoch 80: Train Loss = 0.2251, Train Acc = 0.9141, Val Loss = 0.3234, Val Acc = 0.9062\n",
      "Early stopping at epoch 80\n",
      "Fold 3 - Test Accuracy: 0.7000, Balanced Acc: 0.7179, MCC: 0.4088\n",
      "\n",
      "Fold 4/5\n",
      "Epoch 0: Train Loss = 0.7062, Train Acc = 0.4609, Val Loss = 0.6931, Val Acc = 0.6562\n",
      "Epoch 10: Train Loss = 0.6918, Train Acc = 0.5078, Val Loss = 0.6910, Val Acc = 0.6562\n",
      "Epoch 20: Train Loss = 0.6823, Train Acc = 0.6328, Val Loss = 0.6704, Val Acc = 0.6562\n",
      "Epoch 30: Train Loss = 0.5915, Train Acc = 0.7031, Val Loss = 0.5338, Val Acc = 0.7500\n",
      "Epoch 40: Train Loss = 0.4889, Train Acc = 0.7422, Val Loss = 0.5616, Val Acc = 0.6875\n",
      "Epoch 50: Train Loss = 0.4657, Train Acc = 0.7344, Val Loss = 0.4098, Val Acc = 0.7812\n",
      "Epoch 60: Train Loss = 0.4011, Train Acc = 0.8125, Val Loss = 0.4051, Val Acc = 0.8438\n",
      "Early stopping at epoch 68\n",
      "Fold 4 - Test Accuracy: 0.8250, Balanced Acc: 0.8305, MCC: 0.6319\n",
      "\n",
      "Fold 5/5\n",
      "Epoch 0: Train Loss = 0.7048, Train Acc = 0.4219, Val Loss = 0.6907, Val Acc = 0.3438\n",
      "Epoch 10: Train Loss = 0.6753, Train Acc = 0.6328, Val Loss = 0.6804, Val Acc = 0.7500\n",
      "Epoch 20: Train Loss = 0.5400, Train Acc = 0.7500, Val Loss = 0.5131, Val Acc = 0.7812\n",
      "Epoch 30: Train Loss = 0.4997, Train Acc = 0.7812, Val Loss = 0.4183, Val Acc = 0.7812\n",
      "Epoch 40: Train Loss = 0.4089, Train Acc = 0.8281, Val Loss = 0.3872, Val Acc = 0.8125\n",
      "Epoch 50: Train Loss = 0.3710, Train Acc = 0.8125, Val Loss = 0.3586, Val Acc = 0.8438\n",
      "Epoch 60: Train Loss = 0.3250, Train Acc = 0.8203, Val Loss = 0.3792, Val Acc = 0.8438\n",
      "Early stopping at epoch 67\n",
      "Fold 5 - Test Accuracy: 0.9250, Balanced Acc: 0.9046, MCC: 0.8270\n",
      "\n",
      "Classification Report:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "          -1       0.66      0.85      0.74        67\n",
      "           1       0.91      0.77      0.84       133\n",
      "\n",
      "    accuracy                           0.80       200\n",
      "   macro avg       0.78      0.81      0.79       200\n",
      "weighted avg       0.83      0.80      0.80       200\n",
      "\n",
      "Results and plots saved to ./results/ECG200\n",
      "\n",
      "=== Dataset: ECGFiveDays ===\n",
      "Number of classes: 2\n",
      "Total samples: 884 (Train: 23, Test: 861)\n",
      "Class 1: 442 samples\n",
      "Class 2: 442 samples\n",
      "\n",
      "Fold 1/5\n",
      "Epoch 0: Train Loss = 0.7026, Train Acc = 0.5212, Val Loss = 0.6946, Val Acc = 0.4681\n",
      "Epoch 10: Train Loss = 0.3126, Train Acc = 0.8693, Val Loss = 0.2014, Val Acc = 0.9433\n",
      "Epoch 20: Train Loss = 0.1471, Train Acc = 0.9435, Val Loss = 0.0901, Val Acc = 0.9787\n",
      "Epoch 30: Train Loss = 0.0856, Train Acc = 0.9770, Val Loss = 0.0980, Val Acc = 0.9574\n",
      "Epoch 40: Train Loss = 0.0357, Train Acc = 0.9894, Val Loss = 0.0469, Val Acc = 0.9858\n",
      "Epoch 50: Train Loss = 0.0373, Train Acc = 0.9929, Val Loss = 0.0306, Val Acc = 0.9929\n",
      "Early stopping at epoch 51\n",
      "Fold 1 - Test Accuracy: 0.9887, Balanced Acc: 0.9886, MCC: 0.9776\n",
      "\n",
      "Fold 2/5\n",
      "Epoch 0: Train Loss = 0.6993, Train Acc = 0.4912, Val Loss = 0.6903, Val Acc = 0.6596\n",
      "Epoch 10: Train Loss = 0.3130, Train Acc = 0.8763, Val Loss = 0.2268, Val Acc = 0.9645\n",
      "Epoch 20: Train Loss = 0.1544, Train Acc = 0.9541, Val Loss = 0.0775, Val Acc = 0.9929\n",
      "Epoch 30: Train Loss = 0.0381, Train Acc = 0.9947, Val Loss = 0.0379, Val Acc = 0.9858\n",
      "Epoch 40: Train Loss = 0.0343, Train Acc = 0.9947, Val Loss = 0.0353, Val Acc = 0.9858\n",
      "Epoch 50: Train Loss = 0.0112, Train Acc = 0.9982, Val Loss = 0.0055, Val Acc = 1.0000\n",
      "Epoch 60: Train Loss = 0.0029, Train Acc = 1.0000, Val Loss = 0.0036, Val Acc = 1.0000\n",
      "Early stopping at epoch 67\n",
      "Fold 2 - Test Accuracy: 0.9944, Balanced Acc: 0.9944, MCC: 0.9888\n",
      "\n",
      "Fold 3/5\n",
      "Epoch 0: Train Loss = 0.7028, Train Acc = 0.4700, Val Loss = 0.6847, Val Acc = 0.5532\n",
      "Epoch 10: Train Loss = 0.2056, Train Acc = 0.9258, Val Loss = 0.1739, Val Acc = 0.9362\n",
      "Epoch 20: Train Loss = 0.0804, Train Acc = 0.9788, Val Loss = 0.0538, Val Acc = 0.9929\n",
      "Epoch 30: Train Loss = 0.0394, Train Acc = 0.9929, Val Loss = 0.0294, Val Acc = 0.9929\n",
      "Epoch 40: Train Loss = 0.0286, Train Acc = 0.9929, Val Loss = 0.0688, Val Acc = 0.9787\n",
      "Epoch 50: Train Loss = 0.0698, Train Acc = 0.9753, Val Loss = 0.0380, Val Acc = 0.9929\n",
      "Early stopping at epoch 51\n",
      "Fold 3 - Test Accuracy: 0.9887, Balanced Acc: 0.9888, MCC: 0.9777\n",
      "\n",
      "Fold 4/5\n",
      "Epoch 0: Train Loss = 0.6939, Train Acc = 0.5389, Val Loss = 0.6905, Val Acc = 0.4823\n",
      "Epoch 10: Train Loss = 0.2961, Train Acc = 0.8675, Val Loss = 0.2014, Val Acc = 0.9362\n",
      "Epoch 20: Train Loss = 0.1224, Train Acc = 0.9576, Val Loss = 0.1052, Val Acc = 0.9716\n",
      "Epoch 30: Train Loss = 0.0706, Train Acc = 0.9823, Val Loss = 0.0379, Val Acc = 0.9929\n",
      "Epoch 40: Train Loss = 0.0491, Train Acc = 0.9894, Val Loss = 0.0233, Val Acc = 0.9929\n",
      "Epoch 50: Train Loss = 0.0616, Train Acc = 0.9788, Val Loss = 0.1100, Val Acc = 0.9716\n",
      "Early stopping at epoch 54\n",
      "Fold 4 - Test Accuracy: 0.9944, Balanced Acc: 0.9943, MCC: 0.9888\n",
      "\n",
      "Fold 5/5\n",
      "Epoch 0: Train Loss = 0.6983, Train Acc = 0.5168, Val Loss = 0.6877, Val Acc = 0.5248\n",
      "Epoch 10: Train Loss = 0.3952, Train Acc = 0.8166, Val Loss = 0.3240, Val Acc = 0.8440\n",
      "Epoch 20: Train Loss = 0.3149, Train Acc = 0.8677, Val Loss = 0.2158, Val Acc = 0.9291\n",
      "Epoch 30: Train Loss = 0.1521, Train Acc = 0.9489, Val Loss = 0.1015, Val Acc = 0.9574\n",
      "Epoch 40: Train Loss = 0.0607, Train Acc = 0.9824, Val Loss = 0.0360, Val Acc = 0.9858\n",
      "Epoch 50: Train Loss = 0.0269, Train Acc = 0.9947, Val Loss = 0.0222, Val Acc = 0.9929\n",
      "Epoch 60: Train Loss = 0.0194, Train Acc = 0.9982, Val Loss = 0.0094, Val Acc = 1.0000\n",
      "Epoch 70: Train Loss = 0.0169, Train Acc = 0.9947, Val Loss = 0.0269, Val Acc = 0.9858\n",
      "Early stopping at epoch 71\n",
      "Fold 5 - Test Accuracy: 1.0000, Balanced Acc: 1.0000, MCC: 1.0000\n",
      "\n",
      "Classification Report:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           1       0.99      1.00      0.99       442\n",
      "           2       1.00      0.99      0.99       442\n",
      "\n",
      "    accuracy                           0.99       884\n",
      "   macro avg       0.99      0.99      0.99       884\n",
      "weighted avg       0.99      0.99      0.99       884\n",
      "\n",
      "Results and plots saved to ./results/ECGFiveDays\n",
      "\n",
      "=== Dataset: TwoLeadECG ===\n",
      "Number of classes: 2\n",
      "Total samples: 1162 (Train: 23, Test: 1139)\n",
      "Class 1: 581 samples\n",
      "Class 2: 581 samples\n",
      "\n",
      "Fold 1/5\n",
      "Epoch 0: Train Loss = 0.6888, Train Acc = 0.5255, Val Loss = 0.7167, Val Acc = 0.4703\n",
      "Epoch 10: Train Loss = 0.5445, Train Acc = 0.7406, Val Loss = 0.5493, Val Acc = 0.7027\n",
      "Epoch 20: Train Loss = 0.1832, Train Acc = 0.9382, Val Loss = 0.3301, Val Acc = 0.8919\n",
      "Epoch 30: Train Loss = 0.0958, Train Acc = 0.9704, Val Loss = 0.0499, Val Acc = 0.9892\n",
      "Epoch 40: Train Loss = 0.0610, Train Acc = 0.9772, Val Loss = 0.0282, Val Acc = 0.9892\n",
      "Epoch 50: Train Loss = 0.0538, Train Acc = 0.9839, Val Loss = 0.0583, Val Acc = 0.9838\n",
      "Epoch 60: Train Loss = 0.0297, Train Acc = 0.9933, Val Loss = 0.0480, Val Acc = 0.9838\n",
      "Early stopping at epoch 61\n",
      "Fold 1 - Test Accuracy: 1.0000, Balanced Acc: 1.0000, MCC: 1.0000\n",
      "\n",
      "Fold 2/5\n",
      "Epoch 0: Train Loss = 0.7049, Train Acc = 0.4677, Val Loss = 0.6910, Val Acc = 0.5189\n",
      "Epoch 10: Train Loss = 0.3741, Train Acc = 0.8575, Val Loss = 0.4967, Val Acc = 0.7243\n",
      "Epoch 20: Train Loss = 0.1185, Train Acc = 0.9583, Val Loss = 0.2850, Val Acc = 0.8541\n",
      "Epoch 30: Train Loss = 0.0682, Train Acc = 0.9798, Val Loss = 0.0253, Val Acc = 0.9946\n",
      "Epoch 40: Train Loss = 0.0310, Train Acc = 0.9919, Val Loss = 0.0260, Val Acc = 0.9946\n",
      "Epoch 50: Train Loss = 0.0178, Train Acc = 0.9973, Val Loss = 0.0302, Val Acc = 0.9946\n",
      "Epoch 60: Train Loss = 0.0161, Train Acc = 0.9933, Val Loss = 0.0478, Val Acc = 0.9838\n",
      "Early stopping at epoch 64\n",
      "Fold 2 - Test Accuracy: 0.9957, Balanced Acc: 0.9957, MCC: 0.9915\n",
      "\n",
      "Fold 3/5\n",
      "Epoch 0: Train Loss = 0.7035, Train Acc = 0.4852, Val Loss = 0.7010, Val Acc = 0.4570\n",
      "Epoch 10: Train Loss = 0.3063, Train Acc = 0.8844, Val Loss = 0.1790, Val Acc = 0.9516\n",
      "Epoch 20: Train Loss = 0.0793, Train Acc = 0.9745, Val Loss = 0.0432, Val Acc = 0.9946\n",
      "Epoch 30: Train Loss = 0.0464, Train Acc = 0.9879, Val Loss = 0.1221, Val Acc = 0.9516\n",
      "Epoch 40: Train Loss = 0.0393, Train Acc = 0.9906, Val Loss = 0.0345, Val Acc = 0.9946\n",
      "Epoch 50: Train Loss = 0.0470, Train Acc = 0.9892, Val Loss = 0.0828, Val Acc = 0.9624\n",
      "Epoch 60: Train Loss = 0.0254, Train Acc = 0.9933, Val Loss = 0.0393, Val Acc = 0.9946\n",
      "Early stopping at epoch 65\n",
      "Fold 3 - Test Accuracy: 0.9741, Balanced Acc: 0.9741, MCC: 0.9483\n",
      "\n",
      "Fold 4/5\n",
      "Epoch 0: Train Loss = 0.6968, Train Acc = 0.5013, Val Loss = 0.6893, Val Acc = 0.5161\n",
      "Epoch 10: Train Loss = 0.3001, Train Acc = 0.8911, Val Loss = 0.3235, Val Acc = 0.8817\n",
      "Epoch 20: Train Loss = 0.1173, Train Acc = 0.9624, Val Loss = 0.1714, Val Acc = 0.9462\n",
      "Epoch 30: Train Loss = 0.0710, Train Acc = 0.9758, Val Loss = 0.0499, Val Acc = 0.9839\n",
      "Epoch 40: Train Loss = 0.0413, Train Acc = 0.9879, Val Loss = 0.0357, Val Acc = 0.9946\n",
      "Epoch 50: Train Loss = 0.0664, Train Acc = 0.9839, Val Loss = 0.0258, Val Acc = 0.9946\n",
      "Early stopping at epoch 59\n",
      "Fold 4 - Test Accuracy: 0.9957, Balanced Acc: 0.9957, MCC: 0.9914\n",
      "\n",
      "Fold 5/5\n",
      "Epoch 0: Train Loss = 0.7049, Train Acc = 0.4718, Val Loss = 0.6932, Val Acc = 0.4892\n",
      "Epoch 10: Train Loss = 0.3972, Train Acc = 0.8387, Val Loss = 0.3407, Val Acc = 0.8763\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import copy\n",
    "import json\n",
    "import warnings\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')  # For headless environments\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from scipy.stats import bootstrap\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from sklearn.metrics import (accuracy_score, confusion_matrix, classification_report, roc_curve, auc,\n",
    "                             balanced_accuracy_score, matthews_corrcoef)\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.loader import DataLoader\n",
    "from torch_geometric.nn import GATConv, global_mean_pool\n",
    "from torch_geometric.utils import add_self_loops\n",
    "from tslearn.datasets import UCR_UEA_datasets\n",
    "from captum.attr import Saliency\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# Configuration\n",
    "RANDOM_SEED = 42\n",
    "np.random.seed(RANDOM_SEED)\n",
    "torch.manual_seed(RANDOM_SEED)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed(RANDOM_SEED)\n",
    "\n",
    "MAX_EPOCHS = 100\n",
    "BATCH_SIZE = 32\n",
    "LR = 0.001\n",
    "N_SPLITS = 5\n",
    "HIDDEN_DIM = 64\n",
    "HEADS = 8\n",
    "DROPOUT = 0.2\n",
    "RESULTS_DIR = \"./results\"\n",
    "os.makedirs(RESULTS_DIR, exist_ok=True)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "DATASETS = [\"ECG200\", \"ECGFiveDays\", \"TwoLeadECG\"]  # Example datasets\n",
    "\n",
    "class EarlyStopping:\n",
    "    def __init__(self, patience=10, min_delta=0):\n",
    "        self.patience = patience\n",
    "        self.min_delta = min_delta\n",
    "        self.counter = 0\n",
    "        self.best_loss = None\n",
    "        self.early_stop = False\n",
    "\n",
    "    def __call__(self, val_loss):\n",
    "        if self.best_loss is None:\n",
    "            self.best_loss = val_loss\n",
    "        elif val_loss > self.best_loss - self.min_delta:\n",
    "            self.counter += 1\n",
    "            if self.counter >= self.patience:\n",
    "                self.early_stop = True\n",
    "        else:\n",
    "            self.best_loss = val_loss\n",
    "            self.counter = 0\n",
    "\n",
    "def load_dataset(name):\n",
    "    \"\"\"Load UCR/UEA dataset without normalization (to avoid global data leakage).\n",
    "    We'll do normalization per-fold using training data only.\"\"\"\n",
    "    ucr = UCR_UEA_datasets()\n",
    "    X_train, y_train, X_test, y_test = ucr.load_dataset(name)\n",
    "    \n",
    "    # Squeeze for univariate\n",
    "    if X_train.ndim == 3 and X_train.shape[1] == 1:\n",
    "        X_train = X_train.squeeze(1)\n",
    "    if X_test.ndim == 3 and X_test.shape[1] == 1:\n",
    "        X_test = X_test.squeeze(1)\n",
    "    \n",
    "    unique_labels = np.unique(np.concatenate([y_train, y_test]))\n",
    "    label_to_idx = {lab: i for i, lab in enumerate(unique_labels)}\n",
    "    y_train = np.array([label_to_idx[lab] for lab in y_train], dtype=np.int64)\n",
    "    y_test = np.array([label_to_idx[lab] for lab in y_test], dtype=np.int64)\n",
    "    \n",
    "    X_all = np.concatenate([X_train, X_test], axis=0)\n",
    "    y_all = np.concatenate([y_train, y_test], axis=0)\n",
    "    \n",
    "    return X_all, y_all, unique_labels, (X_train, y_train, X_test, y_test)\n",
    "\n",
    "def normalize_data(X, mean_, std_):\n",
    "    return (X - mean_) / (std_ + 1e-8)\n",
    "\n",
    "def timeseries_to_graph(X, window=5):\n",
    "    \"\"\"Convert time series to graphs: nodes are timesteps; edges connect neighbors within a window.\"\"\"\n",
    "    graphs = []\n",
    "    for i in range(X.shape[0]):\n",
    "        x_val = X[i]\n",
    "        length = x_val.shape[0]\n",
    "        positions = np.linspace(0, 1, length)\n",
    "        x_feat = np.column_stack([x_val, positions])\n",
    "        x_feat = torch.tensor(x_feat, dtype=torch.float32)\n",
    "        \n",
    "        edge_list = []\n",
    "        w = min(window, length)\n",
    "        for j in range(length):\n",
    "            for k in range(max(0, j-w), min(length, j+w+1)):\n",
    "                if j != k:\n",
    "                    edge_list.append([j, k])\n",
    "        \n",
    "        if edge_list:\n",
    "            edge_index = torch.tensor(edge_list, dtype=torch.long).t()\n",
    "        else:\n",
    "            edge_index = torch.empty((2, 0), dtype=torch.long)\n",
    "        \n",
    "        edge_index, _ = add_self_loops(edge_index, num_nodes=length)\n",
    "        data = Data(x=x_feat, edge_index=edge_index)\n",
    "        graphs.append(data)\n",
    "    return graphs\n",
    "\n",
    "class GATConvWithAlpha(GATConv):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.last_alpha = None\n",
    "        self._cached_edge_index = None\n",
    "\n",
    "    def forward(self, x, edge_index, return_attention_weights=False):\n",
    "        out, (edge_index_out, alpha) = super().forward(x, edge_index, return_attention_weights=True)\n",
    "        self.last_alpha = alpha\n",
    "        self._cached_edge_index = edge_index_out\n",
    "        return out\n",
    "\n",
    "class GNNTimeSeriesClassifier(nn.Module):\n",
    "    \"\"\"GAT-based classifier for time series represented as graphs.\"\"\"\n",
    "    def __init__(self, input_dim=2, hidden_dim=64, num_classes=2, heads=8, dropout=0.2):\n",
    "        super().__init__()\n",
    "        self.input_proj = nn.Linear(input_dim, hidden_dim)\n",
    "        self.conv1 = GATConvWithAlpha(hidden_dim, hidden_dim, heads=heads, dropout=dropout)\n",
    "        self.conv2 = GATConvWithAlpha(hidden_dim*heads, hidden_dim, heads=heads, dropout=dropout, concat=False)\n",
    "        \n",
    "        self.ln1 = nn.LayerNorm(hidden_dim * heads)\n",
    "        self.ln2 = nn.LayerNorm(hidden_dim)\n",
    "        \n",
    "        self.fc1 = nn.Linear(hidden_dim, hidden_dim)\n",
    "        self.fc2 = nn.Linear(hidden_dim, num_classes)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "        x = self.input_proj(x)\n",
    "        x = torch.relu(x)\n",
    "        \n",
    "        x1 = self.conv1(x, edge_index)\n",
    "        x1 = self.ln1(x1)\n",
    "        x1 = torch.relu(x1)\n",
    "        x1 = self.dropout(x1)\n",
    "        \n",
    "        x2 = self.conv2(x1, edge_index)\n",
    "        x2 = self.ln2(x2)\n",
    "        x2 = torch.relu(x2)\n",
    "        \n",
    "        x_pool = global_mean_pool(x2, batch)\n",
    "        x_pool = self.fc1(x_pool)\n",
    "        x_pool = torch.relu(x_pool)\n",
    "        x_pool = self.dropout(x_pool)\n",
    "        out = self.fc2(x_pool)\n",
    "        \n",
    "        return out\n",
    "\n",
    "    def get_attention_weights(self, data):\n",
    "        \"\"\"Node-level attention weights from first GAT layer.\"\"\"\n",
    "        self.eval()\n",
    "        with torch.no_grad():\n",
    "            _ = self.forward(data)\n",
    "            alpha1 = self.conv1.last_alpha\n",
    "            edge_idx = self.conv1._cached_edge_index\n",
    "\n",
    "        num_nodes = data.x.size(0)\n",
    "        node_importance = torch.zeros(num_nodes, device=data.x.device)\n",
    "        \n",
    "        alpha_mean = alpha1.mean(dim=1)\n",
    "        for i in range(edge_idx.size(1)):\n",
    "            dst_node = edge_idx[1, i].item()\n",
    "            node_importance[dst_node] += alpha_mean[i].item()\n",
    "        \n",
    "        node_importance = node_importance / (node_importance.sum() + 1e-9)\n",
    "        return node_importance.cpu().numpy()\n",
    "\n",
    "def train_epoch(model, loader, criterion, optimizer, device):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    predictions = []\n",
    "    targets = []\n",
    "    \n",
    "    for batch in loader:\n",
    "        batch = batch.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        output = model(batch)\n",
    "        loss = criterion(output, batch.y)\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        total_loss += loss.item()\n",
    "        predictions.extend(output.argmax(dim=1).cpu().numpy())\n",
    "        targets.extend(batch.y.cpu().numpy())\n",
    "    \n",
    "    return total_loss / len(loader), accuracy_score(targets, predictions)\n",
    "\n",
    "def evaluate(model, loader, criterion, device):\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    predictions = []\n",
    "    targets = []\n",
    "    outputs = []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for batch in loader:\n",
    "            batch = batch.to(device)\n",
    "            output = model(batch)\n",
    "            loss = criterion(output, batch.y)\n",
    "            \n",
    "            total_loss += loss.item()\n",
    "            predictions.extend(output.argmax(dim=1).cpu().numpy())\n",
    "            targets.extend(batch.y.cpu().numpy())\n",
    "            outputs.extend(torch.softmax(output, dim=1).cpu().numpy())\n",
    "    \n",
    "    return (\n",
    "        total_loss / len(loader),\n",
    "        accuracy_score(targets, predictions),\n",
    "        np.array(predictions),\n",
    "        np.array(targets),\n",
    "        np.array(outputs)\n",
    "    )\n",
    "\n",
    "def train_model(model, train_loader, val_loader, criterion, optimizer, device, max_epochs=100):\n",
    "    early_stopping = EarlyStopping(patience=10)\n",
    "    best_val_acc = 0\n",
    "    best_model = None\n",
    "    \n",
    "    for epoch in range(max_epochs):\n",
    "        train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)\n",
    "        val_loss, val_acc, _, _, _ = evaluate(model, val_loader, criterion, device)\n",
    "        \n",
    "        if val_acc > best_val_acc:\n",
    "            best_val_acc = val_acc\n",
    "            best_model = copy.deepcopy(model.state_dict())\n",
    "        \n",
    "        # Logging every 10 epochs\n",
    "        if epoch % 10 == 0:\n",
    "            print(f'Epoch {epoch}: Train Loss = {train_loss:.4f}, Train Acc = {train_acc:.4f}, '\n",
    "                  f'Val Loss = {val_loss:.4f}, Val Acc = {val_acc:.4f}')\n",
    "        \n",
    "        early_stopping(val_loss)\n",
    "        if early_stopping.early_stop:\n",
    "            print(f\"Early stopping at epoch {epoch}\")\n",
    "            break\n",
    "    \n",
    "    model.load_state_dict(best_model)\n",
    "    return model, best_val_acc\n",
    "\n",
    "def bootstrap_confidence_interval(data, stat_func=np.mean, confidence_level=0.95, n_resamples=10000):\n",
    "    \"\"\"Compute bootstrap confidence interval for a given statistic.\"\"\"\n",
    "    res = bootstrap((data,), stat_func, confidence_level=confidence_level, n_resamples=n_resamples, method='basic')\n",
    "    return stat_func(data), res.confidence_interval\n",
    "\n",
    "def plot_confusion_matrix(cm, classes, title, save_path):\n",
    "    plt.figure(figsize=(8, 6))\n",
    "    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=classes, yticklabels=classes)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.ylabel('True Label', fontsize=14)\n",
    "    plt.xlabel('Predicted Label', fontsize=14)\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='png')\n",
    "    plt.close()\n",
    "\n",
    "def plot_roc_with_ci(fprs, tprs, title, save_path):\n",
    "    \"\"\"Plot mean ROC curve with 95% CI.\"\"\"\n",
    "    mean_fpr = np.linspace(0, 1, 100)\n",
    "    interp_tprs = []\n",
    "    for fpr, tpr in zip(fprs, tprs):\n",
    "        interp_tpr = np.interp(mean_fpr, fpr, tpr)\n",
    "        interp_tprs.append(interp_tpr)\n",
    "    interp_tprs = np.array(interp_tprs)\n",
    "    \n",
    "    mean_tpr = interp_tprs.mean(axis=0)\n",
    "    std_tpr = interp_tprs.std(axis=0)\n",
    "    tpr_upper = np.minimum(mean_tpr + 1.96 * std_tpr, 1)\n",
    "    tpr_lower = np.maximum(mean_tpr - 1.96 * std_tpr, 0)\n",
    "    mean_auc = auc(mean_fpr, mean_tpr)\n",
    "    \n",
    "    plt.figure(figsize=(8,6))\n",
    "    plt.plot(mean_fpr, mean_tpr, color='darkorange', lw=2, label=f'Mean ROC (AUC = {mean_auc:.2f})')\n",
    "    plt.fill_between(mean_fpr, tpr_lower, tpr_upper, color='grey', alpha=0.2, label='95% CI')\n",
    "    plt.plot([0,1],[0,1], color='navy', lw=2, linestyle='--')\n",
    "    plt.xlim([0,1])\n",
    "    plt.ylim([0,1.05])\n",
    "    plt.xlabel('False Positive Rate', fontsize=14)\n",
    "    plt.ylabel('True Positive Rate', fontsize=14)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.legend(loc='lower right')\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='png')\n",
    "    plt.close()\n",
    "\n",
    "def visualize_attention(series, attention_weights, title, save_path):\n",
    "    \"\"\"Plot time series with attention weights as a translucent overlay.\"\"\"\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    plt.plot(series, label='Time Series', color='blue')\n",
    "    att_norm = (attention_weights - attention_weights.min()) / (attention_weights.max() - attention_weights.min() + 1e-9)\n",
    "    cmap = plt.cm.Reds\n",
    "    for i in range(len(series)):\n",
    "        color = cmap(att_norm[i])\n",
    "        plt.axvspan(i-0.5, i+0.5, color=color, alpha=0.3)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.xlabel('Time Step', fontsize=14)\n",
    "    plt.ylabel('Normalized Value', fontsize=14)\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='png')\n",
    "    plt.close()\n",
    "\n",
    "def plot_example_timeseries(X, y, unique_labels, dataset_name, save_path):\n",
    "    \"\"\"Plot examples of time series from each class.\"\"\"\n",
    "    plt.figure(figsize=(12, 8))\n",
    "    num_classes = len(unique_labels)\n",
    "    examples_per_class = min(3, len(X)//num_classes)\n",
    "    \n",
    "    for i, cls in enumerate(unique_labels):\n",
    "        class_indices = np.where(y == i)[0]\n",
    "        chosen = np.random.choice(class_indices, size=examples_per_class, replace=False)\n",
    "        for j, idx in enumerate(chosen):\n",
    "            plt.subplot(num_classes, examples_per_class, i*examples_per_class + j + 1)\n",
    "            plt.plot(X[idx], color='blue')\n",
    "            plt.title(f'Class: {cls}', fontsize=12)\n",
    "            plt.xticks([])\n",
    "            plt.yticks([])\n",
    "    plt.suptitle(f'{dataset_name} - Example Time Series', fontsize=16)\n",
    "    plt.tight_layout(rect=[0,0,1,0.95])\n",
    "    plt.savefig(save_path, format='png')\n",
    "    plt.close()\n",
    "\n",
    "def plot_graph_with_attention(data, attention_weights, title, save_path):\n",
    "    \"\"\"Visualize graph structure with nodes colored by attention.\"\"\"\n",
    "    x_coords = data.x[:,1].cpu().numpy() * len(data.x)  # scale position\n",
    "    y_vals = data.x[:,0].cpu().numpy()\n",
    "    att_norm = (attention_weights - attention_weights.min()) / (attention_weights.max() - attention_weights.min() + 1e-9)\n",
    "\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    # Draw edges\n",
    "    edge_index = data.edge_index.cpu().numpy()\n",
    "    for i in range(edge_index.shape[1]):\n",
    "        src, dst = edge_index[:, i]\n",
    "        plt.plot([x_coords[src], x_coords[dst]], [y_vals[src], y_vals[dst]], color='lightgray', linewidth=1, alpha=0.5)\n",
    "\n",
    "    # Draw nodes\n",
    "    sc = plt.scatter(x_coords, y_vals, c=att_norm, cmap='Reds', s=50, edgecolors='black')\n",
    "    plt.colorbar(sc, label='Attention Weight (normalized)')\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.xlabel('Time Step (scaled)')\n",
    "    plt.ylabel('Normalized Value')\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='png')\n",
    "    plt.close()\n",
    "\n",
    "def plot_distribution(data_array, title, xlabel, save_path):\n",
    "    \"\"\"Plot a histogram distribution of given data.\"\"\"\n",
    "    plt.figure(figsize=(8,6))\n",
    "    plt.hist(data_array, bins=30, color='steelblue', edgecolor='black', alpha=0.7)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.xlabel(xlabel)\n",
    "    plt.ylabel('Frequency')\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='png')\n",
    "    plt.close()\n",
    "\n",
    "def plot_mean_ci_profile(all_arrays, title, xlabel, ylabel, save_path):\n",
    "    \"\"\"Plot mean and CI of profiles (e.g. attention over time) across samples.\"\"\"\n",
    "    lengths = [len(a) for a in all_arrays]\n",
    "    min_len = min(lengths)\n",
    "    trimmed = np.array([a[:min_len] for a in all_arrays])\n",
    "    mean_ = trimmed.mean(axis=0)\n",
    "    std_ = trimmed.std(axis=0)\n",
    "\n",
    "    timesteps = np.arange(min_len)\n",
    "    plt.figure(figsize=(10,6))\n",
    "    plt.plot(timesteps, mean_, color='red', linewidth=2, label='Mean')\n",
    "    plt.fill_between(timesteps, mean_ - 1.96*std_, mean_ + 1.96*std_, color='pink', alpha=0.3, label='95% CI')\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.xlabel(xlabel)\n",
    "    plt.ylabel(ylabel)\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_path, format='png')\n",
    "    plt.close()\n",
    "\n",
    "def model_forward_wrapper(model):\n",
    "    def forward(x, edge_index, batch):\n",
    "        data = Data(x=x[0], edge_index=edge_index, batch=batch)\n",
    "        out = model(data)\n",
    "        return out\n",
    "    return forward\n",
    "\n",
    "def compute_saliency(model, data, device):\n",
    "    \"\"\"Compute saliency (gradient-based) on node features.\"\"\"\n",
    "    model.eval()\n",
    "    data = data.to(device)\n",
    "    saliency = Saliency(model_forward_wrapper(model))\n",
    "    # For classification target is a scalar, we get class from data.y\n",
    "    target_class = int(data.y.item())\n",
    "    attributions = saliency.attribute(data.x.unsqueeze(0), target=target_class,\n",
    "                                      additional_forward_args=(data.edge_index, data.batch))\n",
    "    return attributions.squeeze(0).cpu().numpy()\n",
    "\n",
    "def run_experiment(dataset_name):\n",
    "    print(f\"\\n=== Dataset: {dataset_name} ===\")\n",
    "    \n",
    "    # Create dataset-specific directory\n",
    "    ds_dir = os.path.join(RESULTS_DIR, dataset_name)\n",
    "    os.makedirs(ds_dir, exist_ok=True)\n",
    "    \n",
    "    # Load raw data (no normalization)\n",
    "    X_all, y_all, unique_labels, (X_train_orig, y_train_orig, X_test_orig, y_test_orig) = load_dataset(dataset_name)\n",
    "    num_classes = len(unique_labels)\n",
    "\n",
    "    print(f\"Number of classes: {num_classes}\")\n",
    "    print(f\"Total samples: {len(X_all)} (Train: {len(X_train_orig)}, Test: {len(X_test_orig)})\")\n",
    "    class_counts = np.bincount(y_all)\n",
    "    for i, lab in enumerate(unique_labels):\n",
    "        print(f\"Class {lab}: {class_counts[i]} samples\")\n",
    "\n",
    "    # Plot example time series\n",
    "    plot_example_timeseries(X_all, y_all, unique_labels, dataset_name, \n",
    "                            save_path=os.path.join(ds_dir, f\"{dataset_name}_example_ts.png\"))\n",
    "\n",
    "    all_val_accs = []\n",
    "    all_test_accs = []\n",
    "    all_bal_accs = []\n",
    "    all_mccs = []\n",
    "    all_test_preds = []\n",
    "    all_test_targets = []\n",
    "    all_test_probs = []\n",
    "    fprs_list = []\n",
    "    tprs_list = []\n",
    "    all_test_attentions = []\n",
    "    all_test_saliencies = []\n",
    "\n",
    "    skf = StratifiedKFold(n_splits=N_SPLITS, shuffle=True, random_state=RANDOM_SEED)\n",
    "    for fold, (train_idx, test_idx) in enumerate(skf.split(X_all, y_all), 1):\n",
    "        print(f\"\\nFold {fold}/{N_SPLITS}\")\n",
    "\n",
    "        # Split into train/val/test sets for this fold\n",
    "        # val = last 20% of train indices\n",
    "        val_size = int(0.2 * len(train_idx))\n",
    "        train_indices = train_idx[:-val_size]\n",
    "        val_indices = train_idx[-val_size:]\n",
    "        \n",
    "        # Normalize using only train subset\n",
    "        X_train_fold = X_all[train_indices]\n",
    "        mean_ = X_train_fold.mean()\n",
    "        std_ = X_train_fold.std()\n",
    "\n",
    "        X_train_norm = normalize_data(X_all[train_indices], mean_, std_)\n",
    "        X_val_norm = normalize_data(X_all[val_indices], mean_, std_)\n",
    "        X_test_norm = normalize_data(X_all[test_idx], mean_, std_)\n",
    "\n",
    "        # Build graphs\n",
    "        train_graphs = timeseries_to_graph(X_train_norm)\n",
    "        val_graphs = timeseries_to_graph(X_val_norm)\n",
    "        test_graphs = timeseries_to_graph(X_test_norm)\n",
    "        \n",
    "        labels_tensor = torch.tensor(y_all, dtype=torch.long)\n",
    "        \n",
    "        for i, g in enumerate(train_graphs):\n",
    "            g.y = labels_tensor[train_indices[i]].unsqueeze(0)\n",
    "        for i, g in enumerate(val_graphs):\n",
    "            g.y = labels_tensor[val_indices[i]].unsqueeze(0)\n",
    "        for i, g in enumerate(test_graphs):\n",
    "            g.y = labels_tensor[test_idx[i]].unsqueeze(0)\n",
    "\n",
    "        train_loader = DataLoader(train_graphs, batch_size=BATCH_SIZE, shuffle=True)\n",
    "        val_loader = DataLoader(val_graphs, batch_size=BATCH_SIZE, shuffle=False)\n",
    "        test_loader = DataLoader(test_graphs, batch_size=BATCH_SIZE, shuffle=False)\n",
    "        \n",
    "        model = GNNTimeSeriesClassifier(\n",
    "            input_dim=2,\n",
    "            hidden_dim=HIDDEN_DIM,\n",
    "            num_classes=num_classes,\n",
    "            heads=HEADS,\n",
    "            dropout=DROPOUT\n",
    "        ).to(device)\n",
    "        \n",
    "        optimizer = optim.Adam(model.parameters(), lr=LR, weight_decay=1e-4)\n",
    "        \n",
    "        # Class weighting\n",
    "        train_labels_fold = y_all[train_indices]\n",
    "        class_counts_train = np.bincount(train_labels_fold)\n",
    "        class_weights = torch.FloatTensor(1.0 / class_counts_train).to(device)\n",
    "        criterion = nn.CrossEntropyLoss(weight=class_weights)\n",
    "        \n",
    "        model, val_acc = train_model(\n",
    "            model=model,\n",
    "            train_loader=train_loader,\n",
    "            val_loader=val_loader,\n",
    "            criterion=criterion,\n",
    "            optimizer=optimizer,\n",
    "            device=device,\n",
    "            max_epochs=MAX_EPOCHS\n",
    "        )\n",
    "        \n",
    "        test_loss, test_acc, predictions, targets, probabilities = evaluate(\n",
    "            model=model,\n",
    "            loader=test_loader,\n",
    "            criterion=criterion,\n",
    "            device=device\n",
    "        )\n",
    "\n",
    "        bal_acc = balanced_accuracy_score(targets, predictions)\n",
    "        mcc = matthews_corrcoef(targets, predictions)\n",
    "\n",
    "        # Compute AUC if binary\n",
    "        if num_classes == 2:\n",
    "            fpr, tpr, _ = roc_curve(targets, probabilities[:, 1])\n",
    "            fprs_list.append(fpr)\n",
    "            tprs_list.append(tpr)\n",
    "        \n",
    "        # Store metrics\n",
    "        all_val_accs.append(val_acc)\n",
    "        all_test_accs.append(test_acc)\n",
    "        all_bal_accs.append(bal_acc)\n",
    "        all_mccs.append(mcc)\n",
    "        all_test_preds.extend(predictions)\n",
    "        all_test_targets.extend(targets)\n",
    "        all_test_probs.extend(probabilities)\n",
    "\n",
    "        print(f\"Fold {fold} - Test Accuracy: {test_acc:.4f}, Balanced Acc: {bal_acc:.4f}, MCC: {mcc:.4f}\")\n",
    "\n",
    "        # Interpretability on multiple test samples\n",
    "        # Let's pick 3 random test samples to visualize attention and saliency\n",
    "        test_sample_indices = np.random.choice(len(test_graphs), size=min(3, len(test_graphs)), replace=False)\n",
    "        for idx_ in test_sample_indices:\n",
    "            example_data = test_graphs[idx_].clone().to(device)\n",
    "            attention_weights = model.get_attention_weights(example_data)\n",
    "            example_series = example_data.x[:, 0].cpu().numpy()\n",
    "\n",
    "            # Attention on TS\n",
    "            visualize_attention(\n",
    "                series=example_series,\n",
    "                attention_weights=attention_weights,\n",
    "                title=f\"{dataset_name} - Fold {fold} Sample {idx_} Attention (TS)\",\n",
    "                save_path=os.path.join(ds_dir, f\"{dataset_name}_fold{fold}_sample{idx_}_attention_ts.png\")\n",
    "            )\n",
    "\n",
    "            # Graph with attention\n",
    "            plot_graph_with_attention(\n",
    "                data=example_data.cpu(),\n",
    "                attention_weights=attention_weights,\n",
    "                title=f\"{dataset_name} - Fold {fold} Sample {idx_} Graph Attention\",\n",
    "                save_path=os.path.join(ds_dir, f\"{dataset_name}_fold{fold}_sample{idx_}_attention_graph.png\")\n",
    "            )\n",
    "            \n",
    "            # Saliency\n",
    "            saliency_vals = compute_saliency(model, example_data.clone(), device)\n",
    "            # Plot saliency heatmap\n",
    "            plt.figure(figsize=(10,4))\n",
    "            plt.imshow(saliency_vals.T, aspect='auto', cmap='coolwarm')\n",
    "            plt.colorbar(label='Saliency')\n",
    "            plt.title(f\"{dataset_name} - Fold {fold} Sample {idx_} Saliency Heatmap\")\n",
    "            plt.xlabel('Node Index')\n",
    "            plt.ylabel('Feature (0:value, 1:position)')\n",
    "            plt.tight_layout()\n",
    "            plt.savefig(os.path.join(ds_dir, f\"{dataset_name}_fold{fold}_sample{idx_}_saliency.png\"), format='png')\n",
    "            plt.close()\n",
    "\n",
    "        # Gather attention & saliency for all test samples\n",
    "        for d in test_graphs:\n",
    "            d = d.to(device)\n",
    "            aw = model.get_attention_weights(d)\n",
    "            all_test_attentions.append(aw)\n",
    "            sal = compute_saliency(model, d.clone(), device)\n",
    "            # We can aggregate saliency on the \"value\" feature only for interpretability\n",
    "            all_test_saliencies.append(sal[:,0])  # focus on value dimension\n",
    "\n",
    "    # Convert to np arrays\n",
    "    all_test_preds = np.array(all_test_preds)\n",
    "    all_test_targets = np.array(all_test_targets)\n",
    "    all_test_probs = np.array(all_test_probs)\n",
    "\n",
    "    # Compute bootstrap CIs\n",
    "    mean_acc, ci_acc = bootstrap_confidence_interval(np.array(all_test_accs))\n",
    "    mean_bal_acc, ci_bal_acc = bootstrap_confidence_interval(np.array(all_bal_accs))\n",
    "    mean_mcc, ci_mcc = bootstrap_confidence_interval(np.array(all_mccs))\n",
    "\n",
    "    # AUC bootstrap if binary\n",
    "    mean_auc_, low_auc_, high_auc_ = None, None, None\n",
    "    if num_classes == 2:\n",
    "        def auc_stat(data):\n",
    "            targets_, probs_ = data\n",
    "            fpr_, tpr_, _ = roc_curve(targets_.astype(int), probs_)\n",
    "            return auc(fpr_, tpr_)\n",
    "\n",
    "        auc_samples = []\n",
    "        rng = np.random.default_rng(RANDOM_SEED)\n",
    "        for _ in range(10000):\n",
    "            idx = rng.integers(0, len(all_test_targets), len(all_test_targets))\n",
    "            auc_samples.append(auc_stat((all_test_targets[idx], all_test_probs[idx,1])))\n",
    "        auc_samples = np.array(auc_samples)\n",
    "        mean_auc_ = np.mean(auc_samples)\n",
    "        low_auc_ = np.percentile(auc_samples, 2.5)\n",
    "        high_auc_ = np.percentile(auc_samples, 97.5)\n",
    "\n",
    "    # Confusion matrix\n",
    "    cm = confusion_matrix(all_test_targets, all_test_preds)\n",
    "    plot_confusion_matrix(\n",
    "        cm=cm,\n",
    "        classes=[str(l) for l in unique_labels],\n",
    "        title=f\"{dataset_name} Confusion Matrix\",\n",
    "        save_path=os.path.join(ds_dir, f\"{dataset_name}_confusion_matrix.png\")\n",
    "    )\n",
    "\n",
    "    # ROC curve with CI\n",
    "    if num_classes == 2 and len(fprs_list) == N_SPLITS:\n",
    "        plot_roc_with_ci(\n",
    "            fprs=fprs_list,\n",
    "            tprs=tprs_list,\n",
    "            title=f\"{dataset_name} ROC Curve (with 95% CI)\",\n",
    "            save_path=os.path.join(ds_dir, f\"{dataset_name}_roc_curve.png\")\n",
    "        )\n",
    "\n",
    "    # Classification report\n",
    "    cls_report = classification_report(all_test_targets, all_test_preds, target_names=[str(l) for l in unique_labels])\n",
    "    print(\"\\nClassification Report:\")\n",
    "    print(cls_report)\n",
    "\n",
    "    # Plot attention and saliency distributions\n",
    "    all_test_attentions_arr = np.concatenate(all_test_attentions)\n",
    "    plot_distribution(all_test_attentions_arr,\n",
    "                      title=f\"{dataset_name} - Attention Weight Distribution\",\n",
    "                      xlabel='Attention Weight',\n",
    "                      save_path=os.path.join(ds_dir, f\"{dataset_name}_attention_distribution.png\"))\n",
    "\n",
    "    # Average attention profile\n",
    "    plot_mean_ci_profile(all_test_attentions,\n",
    "                         title=f\"{dataset_name} - Average Attention Profile\",\n",
    "                         xlabel='Time Step',\n",
    "                         ylabel='Attention Weight',\n",
    "                         save_path=os.path.join(ds_dir, f\"{dataset_name}_average_attention_profile.png\"))\n",
    "\n",
    "    # Saliency distribution & profile\n",
    "    all_sal = np.concatenate(all_test_saliencies)\n",
    "    plot_distribution(all_sal,\n",
    "                      title=f\"{dataset_name} - Saliency Distribution (Value Feature)\",\n",
    "                      xlabel='Saliency Value',\n",
    "                      save_path=os.path.join(ds_dir, f\"{dataset_name}_saliency_distribution.png\"))\n",
    "\n",
    "    plot_mean_ci_profile(all_test_saliencies,\n",
    "                         title=f\"{dataset_name} - Average Saliency Profile (Value Feature)\",\n",
    "                         xlabel='Time Step',\n",
    "                         ylabel='Saliency',\n",
    "                         save_path=os.path.join(ds_dir, f\"{dataset_name}_average_saliency_profile.png\"))\n",
    "\n",
    "    # Save metrics to CSV and JSON\n",
    "    import pandas as pd\n",
    "\n",
    "    # Per-fold metrics\n",
    "    results_dict = {\n",
    "        'Val_Accuracy': all_val_accs,\n",
    "        'Test_Accuracy': all_test_accs,\n",
    "        'Balanced_Accuracy': all_bal_accs,\n",
    "        'MCC': all_mccs\n",
    "    }\n",
    "    if num_classes == 2:\n",
    "        # Compute per-fold AUC (not bootstrapped, just direct from combined might not be correct per fold,\n",
    "        # but we already have aggregated AUC from all samples)\n",
    "        pass\n",
    "\n",
    "    pd.DataFrame(results_dict).to_csv(os.path.join(ds_dir, f\"{dataset_name}_fold_metrics.csv\"), index=False)\n",
    "\n",
    "    # Summary metrics with CIs\n",
    "    summary_dict = {\n",
    "        'Mean_Accuracy': mean_acc,\n",
    "        'CI_Accuracy_Low': ci_acc.low,\n",
    "        'CI_Accuracy_High': ci_acc.high,\n",
    "        'Mean_Balanced_Accuracy': mean_bal_acc,\n",
    "        'CI_Bal_Acc_Low': ci_bal_acc.low,\n",
    "        'CI_Bal_Acc_High': ci_bal_acc.high,\n",
    "        'Mean_MCC': mean_mcc,\n",
    "        'CI_MCC_Low': ci_mcc.low,\n",
    "        'CI_MCC_High': ci_mcc.high\n",
    "    }\n",
    "\n",
    "    if num_classes == 2 and mean_auc_ is not None:\n",
    "        summary_dict.update({\n",
    "            'Mean_AUC': mean_auc_,\n",
    "            'CI_AUC_Low': low_auc_,\n",
    "            'CI_AUC_High': high_auc_\n",
    "        })\n",
    "\n",
    "    # Save summary to CSV\n",
    "    pd.DataFrame([summary_dict]).to_csv(os.path.join(ds_dir, f\"{dataset_name}_summary_metrics.csv\"), index=False)\n",
    "\n",
    "    # Save summary to JSON\n",
    "    with open(os.path.join(ds_dir, f\"{dataset_name}_summary_metrics.json\"), 'w') as f:\n",
    "        json.dump(summary_dict, f, indent=4)\n",
    "    \n",
    "    # Save classification report\n",
    "    with open(os.path.join(ds_dir, f\"{dataset_name}_classification_report.txt\"), 'w') as f:\n",
    "        f.write(cls_report)\n",
    "\n",
    "    print(f\"Results and plots saved to {ds_dir}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    print(f\"Using device: {device}\")\n",
    "    for dataset_name in DATASETS:\n",
    "        try:\n",
    "            run_experiment(dataset_name)\n",
    "        except Exception as e:\n",
    "            print(f\"Error processing dataset {dataset_name}: {str(e)}\")\n",
    "            continue\n",
    "    \n",
    "    print(\"\\nAll experiments completed. Results saved in:\", RESULTS_DIR)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}