{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "62e7e36e",
   "metadata": {},
   "source": [
    "Requirements: Pytorch, mat73, numpy\n",
    "\n",
    "```pip install mat73```\n",
    "\n",
    "相關論文： https://ieeexplore.ieee.org/document/9124646"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2962af4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "from os.path import *\n",
    "import numpy as np\n",
    "import random\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import time\n",
    "import sys\n",
    "import mat73\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.io import savemat\n",
    "import os\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6151e083",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class DoubleConv(nn.Module):\n",
    "    \"\"\"(convolution => [BN] => ReLU) * 2\"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, mid_channels=None):\n",
    "        super().__init__()\n",
    "        if not mid_channels:\n",
    "            mid_channels = out_channels\n",
    "        self.double_conv = nn.Sequential(\n",
    "            nn.Conv1d(in_channels, mid_channels, kernel_size=3, padding=1),\n",
    "            nn.GroupNorm(num_groups=8, num_channels=mid_channels),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv1d(mid_channels, out_channels, kernel_size=3, padding=1),\n",
    "            nn.GroupNorm(num_groups=8, num_channels=out_channels),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.double_conv(x)\n",
    "\n",
    "\n",
    "class Down(nn.Module):\n",
    "    \"\"\"Downscaling with maxpool then double conv\"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "        self.maxpool_conv = nn.Sequential(\n",
    "            nn.MaxPool1d(2),\n",
    "            DoubleConv(in_channels, out_channels)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.maxpool_conv(x)\n",
    "\n",
    "\n",
    "class Up(nn.Module):\n",
    "    \"\"\"Upscaling then double conv\"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "\n",
    "        self.up = nn.Upsample(\n",
    "            scale_factor=2, mode='linear', align_corners=True)\n",
    "        self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)\n",
    "\n",
    "    \n",
    "    def forward(self, x1, x2):\n",
    "        x1 = self.up(x1)\n",
    "        # input is CHW\n",
    "        diffX = x2.size()[2] - x1.size()[2]\n",
    "\n",
    "        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2])\n",
    "        # if you have padding issues, see\n",
    "        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a\n",
    "        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd\n",
    "        x = torch.cat([x2, x1], dim=1)\n",
    "        return self.conv(x)\n",
    "\n",
    "\n",
    "class OutConv(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(OutConv, self).__init__()\n",
    "        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "\n",
    "class UNet1d(nn.Module):\n",
    "    def __init__(self, n_channels, n_classes, nfilter=24):\n",
    "        super(UNet1d, self).__init__()\n",
    "        self.n_channels = n_channels\n",
    "        self.n_classes = n_classes\n",
    "\n",
    "        self.inc = DoubleConv(n_channels, nfilter)\n",
    "        self.down1 = Down(nfilter, nfilter * 2)\n",
    "        self.down2 = Down(nfilter * 2, nfilter * 4)\n",
    "        self.down3 = Down(nfilter * 4, nfilter * 8)\n",
    "        self.down4 = Down(nfilter * 8, nfilter * 8)\n",
    "        self.up1 = Up(nfilter * 16, nfilter * 4)\n",
    "        self.up2 = Up(nfilter * 8, nfilter * 2)\n",
    "        self.up3 = Up(nfilter * 4, nfilter * 1)\n",
    "        self.up4 = Up(nfilter * 2, nfilter)\n",
    "        self.outc = OutConv(nfilter, n_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x1 = self.inc(x)\n",
    "        x2 = self.down1(x1)\n",
    "        x3 = self.down2(x2)\n",
    "        x4 = self.down3(x3)\n",
    "        x5 = self.down4(x4)\n",
    "        x = self.up1(x5, x4)\n",
    "        x = self.up2(x, x3)\n",
    "        x = self.up3(x, x2)\n",
    "        x = self.up4(x, x1)\n",
    "        logits = self.outc(x)\n",
    "        return logits\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f21a6a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.signal import butter, sosfilt, sosfreqz\n",
    "\n",
    "def butter_bandpass(lowcut, highcut, fs, order=5):\n",
    "        nyq = 0.5 * fs\n",
    "        low = lowcut / nyq\n",
    "        high = highcut / nyq\n",
    "        sos = butter(order, [low, high], analog=False, btype='band', output='sos')\n",
    "        return sos\n",
    "\n",
    "def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):\n",
    "        sos = butter_bandpass(lowcut, highcut, fs, order=order)\n",
    "        y = sosfilt(sos, data)\n",
    "        return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71b4e84a",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = r'/NFS/tyhuang/fhlinXBCG/noscan/170320_CLY/analysis/EyeClose1_noscan.mat'\n",
    "\n",
    "def norm_ecg(ecg):\n",
    "    min1, max1 = np.percentile(ecg, [1, 99])\n",
    "    ecg[ecg>max1] = max1\n",
    "    ecg[ecg<min1] = min1\n",
    "    ecg = (ecg - min1)/(max1-min1)\n",
    "    return ecg\n",
    "\n",
    "\n",
    "data = mat73.loadmat(f)\n",
    "eeg_filtered = data['EEG_before_bcg'] * 0\n",
    "t = time.time()\n",
    "for ii in range(31):\n",
    "    eeg_filtered[ii, ...] = butter_bandpass_filter(data['EEG_before_bcg'][ii,:], 1, 40, 5000)\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "device = ('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "NET = UNet1d(n_channels=1, n_classes=31, nfilter=8).to(device)\n",
    "#NET = torch.load('pretrainbcg_f8.pt')\n",
    "#NET.outc = OutConv(8, 31)\n",
    "NET = NET.to(device)\n",
    "optimizer = torch.optim.Adam(NET.parameters(), lr=5e-4)\n",
    "optimizer.zero_grad()\n",
    "maxlen = data['ECG'].size\n",
    "\n",
    "loss_list = []\n",
    "count = 0\n",
    "ecg = norm_ecg(data['ECG'])\n",
    "for ii in range(5000):\n",
    "    if ii % 10 == 0:\n",
    "        sys.stdout.write('.')\n",
    "    index = random.randrange(maxlen-20000)\n",
    "    ECG = ecg[index:(index+20000)]\n",
    "    EEG = eeg_filtered[:, index:(index+20000)]\n",
    "    ECG_d = torch.from_numpy(ECG[None, ...][None, ...]).to(device).float()\n",
    "    EEG_d = torch.from_numpy(EEG[None, ...]).to(device).float()\n",
    "\n",
    "    # step 3: forward path of UNET\n",
    "    logits = NET(ECG_d)\n",
    "    loss = nn.MSELoss()(logits, EEG_d)\n",
    "    loss_list.append(loss.item())\n",
    "\n",
    "\n",
    "    # Step 5: Perform back-propagation\n",
    "    loss.backward() #accumulate the gradients\n",
    "    optimizer.step() #Update network weights according to the optimizer\n",
    "    optimizer.zero_grad() #empty the gradients\n",
    "\n",
    "\n",
    "    if (ii + 1) % 500 == 0: #plot results per 500 iterations\n",
    "        print('mse loss: ', np.mean(loss_list))\n",
    "        loss_list = []\n",
    "        EEG = eeg_filtered[:, :50000]\n",
    "        ECG = data['ECG'][:50000]\n",
    "        ECG_d = torch.from_numpy(ECG[None, ...][None, ...]).to(device).float()\n",
    "        EEG_d = torch.from_numpy(EEG[None, ...]).to(device).float()\n",
    "        logits = NET(ECG_d)        \n",
    "        EEG_pred = logits.cpu().detach().numpy()\n",
    "        plt.figure(figsize=(12, 6), dpi=300)\n",
    "\n",
    "        plt.plot(EEG[0, ...], 'g')\n",
    "        plt.plot(EEG[0, ...] - EEG_pred[0, 0, ...], 'r')\n",
    "        time1 = round(time.time() - t, 1)\n",
    "        plt.title(f' {time1} seconds')\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "# remove BCG from the whole dataset\n",
    "EEG = eeg_filtered\n",
    "ECG = data['ECG']\n",
    "ECG_d = torch.from_numpy(ECG[None, ...][None, ...]).to(device).float()\n",
    "EEG_d = torch.from_numpy(EEG[None, ...]).to(device).float()\n",
    "logits = NET(ECG_d)\n",
    "BCG_pred = logits.cpu().detach().numpy()[0, ...]\n",
    "EEG_removeBCG_unet = eeg_filtered - BCG_pred\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dc44936",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}