{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from matplotlib import pylab as plt\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import balanced_accuracy_score\n",
    "\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "\n",
    "from core.io import filter_events, get_ranges, split_epochs_binary\n",
    "from core.signal_processing import filter_signal, __create_constants, CSPFilter\n",
    "from core.nn.provider import NetProvider\n",
    "from core.nn.data import TableDataset\n",
    "\n",
    "%matplotlib inline\n",
    "plt.style.use('bmh')\n",
    "plt.rcParams.update({'font.size': 14})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Необходимо проверить точность классификации с помощью полносвязной сети:\n",
    "\n",
    "1. Подобрать архитектуру (поиск по сетке) - возможно полносвязные слои заменить на свёртку\n",
    "2. Выбрать электроды\n",
    "3. Исследовать влияние фильтрации\n",
    "4. Исследовать влияние использования CSP-фильтра\n",
    "\n",
    "> По числовым значениям - **обсудить с Андреем Николаевичем**\n",
    "\n",
    "> **Окружение TacticENV**\n",
    "\n",
    "> **Сохранять все результаты - ребятам понадабятся графики/таблицы в статье, необзодимые визуализации обсудить с Андреем Николаевичем**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Примеры использования библеотеки `core`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Данные"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Загрузка"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "openvibe_config = json.load(open(\"./config/openVibe.json\", \"r\"))\n",
    "raw_data = pd.read_csv(\"./Записи/OpenVibe signals/motor-imagery-csp-1-acquisition-[2024.04.18-10.37.21].csv\")\n",
    "data = filter_events(raw_data, openvibe_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Фильтрация"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "constants = __create_constants(low=8.0, high=30.0, sample_rate=250.0, order=5) \n",
    "filtered_data = filter_signal(df=data, cols=[x for x in data.columns if \"Channel\" in x], constants=constants)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Получение отрезков и разметки"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = split_epochs_binary(\n",
    "    df=filtered_data, \n",
    "    epoch_duration=4 * 250, \n",
    "    init_pause=int(0.5 * 250), \n",
    "    final_pause=0, \n",
    "    step=50000, # берем только 1 отрезок на действие (обсудить с А.Н.), подробнее см. доку \n",
    "    map_labels={\"left\": 0, \"right\": 1},\n",
    "    cols=[x for x in data.columns if \"Channel\" in x]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CSP-фильтр"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(40, 1000)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "csp_filter = CSPFilter(4)\n",
    "csp_filter.fit(x, y)\n",
    "csp_filter.transform(x).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Нейронки\n",
    "\n",
    "> Используем `pytorch-lightning`, если будешь добавлять - придерживайся стиля 🤖\n",
    "\n",
    "Ниже пример полного пайплайна обучения (без фильтрации, без сколбзящего окна, без CSP-фильтра) для 3 случайных электродов - `Channel 1`, `Channel 2`, `Channel 3`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(40, 3000) (40, 2)\n"
     ]
    }
   ],
   "source": [
    "raw_data = pd.read_csv(\"./Записи/OpenVibe signals/motor-imagery-csp-1-acquisition-[2024.04.18-10.37.21].csv\")\n",
    "data = filter_events(raw_data, openvibe_config)\n",
    "x, y = split_epochs_binary(\n",
    "    df=data, \n",
    "    epoch_duration=4 * 250, \n",
    "    init_pause=int(0.5 * 250), \n",
    "    final_pause=0, \n",
    "    step=50000, # берем только 1 отрезок на действие (обсудить с А.Н.), подробнее см. доку \n",
    "    map_labels={\"left\": 0, \"right\": 1},\n",
    "    cols=[\"Channel 1\", \"Channel 2\", \"Channel 3\"]\n",
    ")\n",
    "\n",
    "# Делаем сигналы плоскими, лейблы в OneHot\n",
    "x = x.reshape((-1, 3 * 1000))\n",
    "y =  np.eye(2)[y]\n",
    "print(x.shape, y.shape)\n",
    "\n",
    "# бъем train/test\n",
    "X_train, X_test, y_train, y_test = train_test_split(x, y, train_size=0.8)\n",
    "\n",
    "# конвертим в объекты торча\n",
    "train_loader = torch.utils.data.DataLoader(TableDataset(X_train, y_train), batch_size=32)\n",
    "val_loader = torch.utils.data.DataLoader(TableDataset(X_test, y_test), batch_size=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = torch.nn.Sequential(\n",
    "    torch.nn.Linear(3000, 1024),\n",
    "    torch.nn.Sigmoid(),\n",
    "    torch.nn.Linear(1024, 2),\n",
    "    torch.nn.Softmax(dim=1),\n",
    ")\n",
    "\n",
    "wrapped_net = NetProvider(\n",
    "    net=net, \n",
    "    lr=1e-4, \n",
    "    criteria=torch.nn.BCELoss(), \n",
    "    metrics=[\n",
    "        lambda pred, labels: balanced_accuracy_score(np.argmax(labels, axis=1), np.argmax(pred, axis=1))\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "Missing logger folder: runs\\example\\lightning_logs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name     | Type       | Params\n",
      "----------------------------------------\n",
      "0 | criteria | BCELoss    | 0     \n",
      "1 | net      | Sequential | 3.1 M \n",
      "----------------------------------------\n",
      "3.1 M     Trainable params\n",
      "0         Non-trainable params\n",
      "3.1 M     Total params\n",
      "12.300    Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc2a3deccc5b4690a005b7309e3f9322",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\user\\anaconda3\\envs\\TacticENV\\lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:441: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n",
      "c:\\Users\\user\\anaconda3\\envs\\TacticENV\\lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:441: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "97215ced575c46d99b8e3cd20ab809e9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d57332f35bf49ab8ff89e3f4215180b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "64ab7eb3058c4e518478349348fcb950",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c621ba6368341d68765f570aaa2b8db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4fc7906c46b44fc8814452719f69a6cb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dcdd8be5bf65421db1af4cf9303da835",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb3f5bc1991b4addb12cd99b9213a57b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "734438becdbc443e87942898e64199a4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0baf18f624a341efba714be21cf42c5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d39784d032f84f5cbc06b1694beb1acc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b1998c91254e493893112c63a0e0b680",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=10` reached.\n"
     ]
    }
   ],
   "source": [
    "torch.set_float32_matmul_precision('medium')\n",
    "trainer = pl.Trainer(\n",
    "    default_root_dir=\"./runs/example\",\n",
    "    accelerator=\"auto\",\n",
    "    max_epochs=10,\n",
    "    callbacks=[\n",
    "        ModelCheckpoint(save_weights_only=True, monitor=\"val/loss\", mode=\"min\"),\n",
    "    ],\n",
    "    log_every_n_steps=1,\n",
    ")\n",
    "\n",
    "trainer.fit(wrapped_net, train_loader, val_loader)\n",
    "__best_path = list(Path(\"./runs/example/lightning_logs/version_0/checkpoints/\").glob(\"*\"))[0]\n",
    "wrapped_net.load_state_dict(torch.load(str(__best_path))['state_dict'])\n",
    "del __best_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
